#208MediumTries

Implement Trie (Prefix Tree)

The tell: Prefix queries over a dictionary, where lookup should cost word length not dictionary size.

How you get there

1Each node holds a map from character to child node.
2A terminal flag distinguishes a stored word from a mere prefix.
3Insert, search and startsWith are all the same walk, differing only in what they check at the end.

Solution

Python

class TrieNode:
    def __init__(self) -> None:
        self.children: dict[str, TrieNode] = {}
        self.is_word = False

class Trie:
    def __init__(self) -> None:
        self.root = TrieNode()

    def insert(self, word: str) -> None:
        node = self.root
        for ch in word:
            node = node.children.setdefault(ch, TrieNode())
        node.is_word = True

    def _walk(self, prefix: str):
        node = self.root
        for ch in prefix:
            if ch not in node.children:
                return None
            node = node.children[ch]
        return node

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

    def startsWith(self, prefix: str) -> bool:
        return self._walk(prefix) is not None

Time: O(len(word)) per operation
Space: O(total characters)

What goes wrong

Without the is_word flag, searching for a stored prefix wrongly reports a full word.

Test yourself on DSA Next: Design Add and Search Words Data Structure →