LC 208 - Implement Trie (Prefix Tree)

Question

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.

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:

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.

Links

Question here and solution here

Solution

concept

Trie basically is a tree which each node represent a char. We have to define a TrieNode and we can use the self.children hashmap to keep track its children node. To add an additional children node, we just need to create a TrieNode and add to the key of that char. Note that the TrieNode itself do not contain information the char but it is contained in the key of the hashmap.

code

class TrieNode:
    def __init__(self):
        self.children = {} # new node will be self.children[c] = TrieNode()
        self.end_of_word = False

class Trie:

    def __init__(self):
        self.root = TrieNode()

    def insert(self, word: str) -> None:
        cur = self.root
        for c in word:
            if c not in cur.children:
                cur.children[c] = TrieNode()
            cur = cur.children[c]
        cur.end_of_word = True
        
    def search(self, word: str) -> bool:
        cur = self.root
        for c in word:
            if c not in cur.children:
                return False
            cur = cur.children[c]
        return cur.end_of_word
        
    def startsWith(self, prefix: str) -> bool:
        cur = self.root
        for c in prefix:
            if c not in cur.children:
                return False
            cur = cur.children[c]
        return True
        


# Your Trie object will be instantiated and called as such:
# obj = Trie()
# obj.insert(word)
# param_2 = obj.search(word)
# param_3 = obj.startsWith(prefix)

Complexity

time: $O(n)$
space: $O(t)$
where n is the length of the string and t is the total number of TrieNode created in the Trie.