LeetCode 208: Implement Trie (Prefix Tree)

link

We represent the Trie as a tree with TrieNode‘s. An edge has two end points and a label: $\texttt{from} \stackrel{\text{label}}{\longrightarrow} \texttt{to}$ .

class TrieNode:
    def __init__(self):
        self.end_of_word = False
        self.edges = {}

Following an edge, $\textcircled{u} \stackrel{c}{\longrightarrow} \textcircled{v}$ , consumes the word character, $c$ . If a node is reached consuming the last character of a word, that node is marked with end_of_word = True.

The root is a special node, it is reached without consuming a word character. As a consequence, "" is a prefix of all words. If we insert "", root will have end_of_word = True.

Recursive

If longest word length is $m$ :

class TrieNode:
    def __init__(self):
        self.end_of_word = False
        self.edges = {}

class Trie:

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

    def __insert(self, u, word, i) -> None:
        if i == len(word):
            u.end_of_word = True
            return
        
        char = word[i]
        if char not in u.edges:
            u.edges[char] = TrieNode()
        self.__insert(u.edges[char], word, i+1)

    def insert(self, word: str) -> None:
        self.__insert(self.root, word, 0)

    def __get(self, u, word, i) -> Optional[TrieNode]:
        if i == len(word):
            return u
        
        char = word[i]
        if char not in u.edges:
            return None
        return self.__get(u.edges[char], word, i+1)

    def search(self, word: str) -> bool:
        node = self.__get(self.root, word, 0)
        return node is not None and node.end_of_word

    def startsWith(self, prefix: str) -> bool:
        node = self.__get(self.root, prefix, 0)
        if node is None:
            return False
        return node.end_of_word or len(node.edges) > 0


# 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)

Iterative

class TrieNode:
    def __init__(self):
        self.end_of_word = False
        self.edges = {}

class Trie:

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

    def insert(self, word: str) -> None:
        u = self.root
        for char in word:
            if char not in u.edges:
                u.edges[char] = TrieNode()
            u = u.edges[char]
        u.end_of_word = True

    def __get(self, u, word) -> Optional[TrieNode]:
        for char in word:
            if char not in u.edges:
                return None
            u = u.edges[char]
        return u

    def search(self, word: str) -> bool:
        node = self.__get(self.root, word)
        return node is not None and node.end_of_word

    def startsWith(self, prefix: str) -> bool:
        node = self.__get(self.root, prefix)
        if node is None:
            return False
        return node.end_of_word or len(node.edges) > 0


# 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)

Operation	Time	Space
`insert`	$\mathcal{O}(m)$	$\mathcal{O}(m)$
`search`	$\mathcal{O}(m)$	$\mathcal{O}(m)$
`startsWith`	$\mathcal{O}(m)$	$\mathcal{O}(m)$

Operation	Time	Space
`insert`	$\mathcal{O}(m)$	$\mathcal{O}(m)$
`search`	$\mathcal{O}(m)$	$\mathcal{O}(1)$
`startsWith`	$\mathcal{O}(m)$	$\mathcal{O}(1)$

unreasonably effective