Implement Trie (Prefix Tree)

Interview Prep

Warm-upTriedesignstringstrie

The problem

Build a trie (prefix tree) supporting three operations: insert(word), search(word) (exact match), and startsWith(prefix). All three should be O(L) where L is the string length, with no dependence on the number of stored words.

Pattern: per-character branching with terminal flag

Each node is a small map from a character to a child node, plus a boolean marking whether some inserted word ends exactly here. Walking the tree consumes one character per step; the difference between search and startsWith is just the terminal-flag check at the end. Two functions, one walk.

The terminal flag is what lets the trie distinguish stored words from mere prefixes. Without it, after inserting "apple" the trie wouldn't be able to say whether "app" is a stored word or just a path on the way somewhere.

Solution

class Trie:
    """Prefix tree over arbitrary alphabet. Each node holds a children dict + a terminal flag."""
    def __init__(self):
        self.children: dict[str, 'Trie'] = {}
        self.is_word: bool = False

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

    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

    def _walk(self, s: str) -> 'Trie | None':
        node = self
        for ch in s:
            node = node.children.get(ch)
            if node is None:
                return None
        return node

Trace

Insert "app", "apple", "ape", "bat":

(root)
├── 'a' -> A
│   ├── 'p' -> AP
│   │   ├── 'p' -> APP (is_word)
│   │   │   └── 'l' -> APPL
│   │   │       └── 'e' -> APPLE (is_word)
│   │   └── 'e' -> APE (is_word)
└── 'b' -> B
    └── 'a' -> BA
        └── 't' -> BAT (is_word)

search("app")   : walk a -> p -> p, node.is_word = True  -> True
search("appl")  : walk a -> p -> p -> l, node.is_word = False -> False
search("apple") : walk a -> p -> p -> l -> e, True -> True
startsWith("ap"): walk a -> p succeeds -> True
startsWith("ax"): walk a -> x, child missing -> False

Variant: fixed-size array for ASCII

If the alphabet is small and fixed (e.g., 26 lowercase letters), replace the dict with a 26-slot array. Slightly faster lookups, slightly more memory for sparse nodes. Choose based on alphabet size and density. 

class TrieArray:
    """Lower-cased ASCII only — slightly faster, slightly less memory."""
    __slots__ = ('next', 'is_word')
    def __init__(self):
        self.next: list['TrieArray | None'] = [None] * 26
        self.is_word: bool = False
    def insert(self, word):
        node = self
        for ch in word:
            i = ord(ch) - ord('a')
            if node.next[i] is None:
                node.next[i] = TrieArray()
            node = node.next[i]
        node.is_word = True

Complexity

Variations worth knowing