Trie (Prefix Tree)

Posted on Edited on In , Views:

What is Trie

A trie (pronounced as “try”) or prefix tree is a tree data structure used to efficiently store and retrieve keys in a dataset of strings. There are various applications of this data structure, such as autocomplete and spellchecker.

[Leetcode 208] Implement Trie (Prefix Tree)

Implement the Trie class:

  • Trie() Initializes the trie object.
  • void insert(String word) Inserts the string word into the trie.
  • boolean search(String word) Returns true if the string word is in the trie (i.e., was inserted before), and false otherwise.
  • boolean startsWith(String prefix) Returns true if there is a previously inserted string word that has the prefix prefix, and false otherwise.

Example 1:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
Input
["Trie", "insert", "search", "search", "startsWith", "insert", "search"]
[[], ["apple"], ["apple"], ["app"], ["app"], ["app"], ["app"]]
Output
[null, null, true, false, true, null, true]

Explanation
Trie trie = new Trie();
trie.insert("apple");
trie.search("apple");   // return True
trie.search("app");     // return False
trie.startsWith("app"); // return True
trie.insert("app");
trie.search("app");     // return True

Constraints:

  • 1 <= word.length, prefix.length <= 2000
  • word and prefix consist only of lowercase English letters.
  • At most 3 * 104 calls in total will be made to insert, search, and startsWith.

Solution:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
class Trie:
    def __init__(self):
        self.root = {}

    # e.g. {'a': {'p': {'p': {'l': {'e': {'*': ''}}}}}}
    def insert(self, word: str) -> None:
        curr = self.root
        for letter in word:
            if letter not in curr:
                curr[letter] = {}
            curr = curr[letter]
        curr['*'] = ''

    def search(self, word: str) -> bool:
        curr = self.root
        for letter in word:
            if letter not in curr:
                return False
            curr = curr[letter]

        return '*' in curr

    def startsWith(self, prefix: str) -> bool:
        curr = self.root
        for letter in prefix:
            if letter not in curr:
                return False
            curr = curr[letter]

        return True

[Leetcode 421] Maximum XOR of Two Numbers in an Array

Given an integer array nums, return the maximum result of nums[i] XOR nums[j], where 0 <= i <= j < n.

Example 1:

1
2
3
4
5
6
7
Input: nums = [3,10,5,25,2,8]
Output: 28
Explanation: The maximum result is 5 XOR 25 = 28.
Example 2:

Input: nums = [14,70,53,83,49,91,36,80,92,51,66,70]
Output: 127

Constraints:

  • 1 <= nums.length <= 2 * 105
  • 0 <= nums[i] <= 231 - 1

Solution:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
class Trie:
    def __init__(self, n):
        self.root = {}
        self.maxLen = n  # max length of all numbers                             
        
    def add_num(self, num):
        node = self.root 
        # shift self.maxLen bits 
        for shift in range(self.maxLen, -1, -1):      
            # get a bit from left to right
            val = 1 if num & (1 << shift) else 0

            if val not in node:
                node[val] = {}

            node = node[val]
        # mark the number
        node['*'] = num
        
class Solution:
    def findMaximumXOR(self, nums: List[int]) -> int:
        # get max length of all numbers' binary. e.g. bin(6)=0b110
        max_len = len(bin(max(nums))) - 2

        # build trie with number's binary
        trie = Trie(max_len)
        for num in nums: trie.add_num(num)      
        #print(trie.root)

        ans = 0

        # for each num, find the number which can create max value with num using XOR
        for num in nums:                              
            node = trie.root 
            for shift in range(max_len, -1, -1):
                # get a bit from left to right
                val = 1 if num & (1 << shift) else 0  

                # try opposite bit if it's in trie to make it larger, otherwise use itself
                node = node[1-val] if 1-val in node else node[val] 

            # calculate xor and save to answer
            ans = max(ans, num ^ node['*'])
        return ans

[Leetcode 1268] Search Suggestions System

You are given an array of strings products and a string searchWord.

Design a system that suggests at most three product names from products after each character of searchWord is typed. Suggested products should have common prefix with searchWord. If there are more than three products with a common prefix return the three lexicographically minimums products.

Return a list of lists of the suggested products after each character of searchWord is typed.

Example 1:

1
2
3
4
5
Input: products = ["mobile","mouse","moneypot","monitor","mousepad"], searchWord = "mouse"
Output: [["mobile","moneypot","monitor"],["mobile","moneypot","monitor"],["mouse","mousepad"],["mouse","mousepad"],["mouse","mousepad"]]
Explanation: products sorted lexicographically = ["mobile","moneypot","monitor","mouse","mousepad"].
After typing m and mo all products match and we show user ["mobile","moneypot","monitor"].
After typing mou, mous and mouse the system suggests ["mouse","mousepad"].

Example 2:

1
2
3
Input: products = ["havana"], searchWord = "havana"
Output: [["havana"],["havana"],["havana"],["havana"],["havana"],["havana"]]
Explanation: The only word "havana" will be always suggested while typing the search word.

Constraints:

  • 1 <= products.length <= 1000
  • 1 <= products[i].length <= 3000
  • 1 <= sum(products[i].length) <= 2 * 104
  • All the strings of products are unique.
  • products[i] consists of lowercase English letters.
  • 1 <= searchWord.length <= 1000
  • searchWord consists of lowercase English letters.

Solution:

Brute force:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
class Solution:
    def suggestedProducts(self, products: List[str], searchWord: str) -> List[List[str]]:
        res = []
        products.sort() # ensure the lexicographical order of products

        for i, c in enumerate(searchWord):
            curr = []
            prefix = searchWord[:i+1]
            for prod in products:
                if prefix == prod[:i+1]:
                    curr.append(prod)
                if len(curr) == 3:
                    break

            res.append(curr)
        
        return res

Using Trie:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
class Trie:
    def __init__(self):
        self.root = {}

    def add(self, word):
        node = self.root
        for c in word:
            if c not in node:
                node[c] = {} # create a child node {}
            
            node = node[c] # move down to the child node

            # add word list to current node
            if 'words' not in node:
                node['words'] = []

            # add word to list since it goes through the current character
            if len(node['words']) < 3:
                node['words'].append(word)          

    def search(self, prefix):
        node = self.root
        for c in prefix:
            if c not in node:
                return ''
            node = node[c]

        return node['words']

class Solution:
    def suggestedProducts(self, products: List[str], searchWord: str) -> List[List[str]]:
        products.sort()

        trie = Trie()
        for prod in products:
            trie.add(prod)

        res = []
        prefix = ""
        for c in searchWord:
            prefix += c
            res.append(trie.search(prefix))

        return res