Trie (Prefix Tree)

Author

Pham Nguyen Hung

Published

July 21, 2023

Trie

Definitions

Problems

1. Implement Trie (Prefix Tree)

Intuition

The first thing I think of is a linked list. Each node containing a character and a pointer to the next node. It was insufficient - I also need to know if the node is the end of a word or not. Okay, I can add a variable. But then there are words sharing the same prefix such as “app” and “apple”. The second “p” is the end, but it also needs to contain a pointer to “l”. And one pointer may not be enough as there may be another word such as “append”. So it is best if I have a dictionary of next nodes. Tada!

And the dictionary can serve another purpose: it can be used to store the value at the node. Explaining in words will be clunky, so here’s a (wrong) illustration.

Each level now holds a dictionary, containing up to 26 lowercase English letters per the problem description. Each letter is the key to a node in the next level, which stores if it is the end of a word and a dictionary connected to the next level. And so on.

All methods in the class share the same crux: iterate the input string and iterate the level.

Complexity

  • Time complexity: \(O(n)\) - for all operations, as the algorithms iterate the input string.

  • Space complexity: \(O(1)\) - ignoring the information stored in each node, each method takes \(O(1)\) space as only pointers are stored.

Code

class _TrieNode:
    def __init__(self) -> None:
        self.children = {}
        self.isEndWord = False

class Trie:
    def __init__(self):
        self.root = _TrieNode()

    def insert(self, word: str) -> None:
        curr = self.root
        
        for char in word:
            if char not in curr.children:
                curr.children[char] = _TrieNode()
            curr = curr.children[char]
        
        curr.isEndWord = True

    def search(self, word: str) -> bool:
        curr = self.root

        for char in word:
            if char not in curr.children:
                return False
            curr = curr.children[char]
        
        return curr.isEndWord

    def startsWith(self, prefix: str) -> bool:
        curr = self.root

        for char in prefix:
            if char not in curr.children:
                return False
            curr = curr.children[char]
        
        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)