Implement a Trie

This is the warm-up trie question, and it is asked often. You are graded on two things: writing the structure without hesitation, and stating the complexity before you are asked. It is also the base for every other problem in this section, so it is worth knowing cold.

The problem

Design a class that stores words and supports three operations:

• insert(word) — store the word.
• search(word) — return True if exactly this word was inserted.
• starts_with(prefix) — return True if any inserted word begins with prefix.

Example. Insert car, cart, cat and dog. Then:

• search("car") → True. It was inserted.
• search("ca") → False. It is only the start of other words.
• starts_with("ca") → True. car, cart and cat all begin with it.
• starts_with("cab") → False. No word continues ca with b.

Constraints. Words and prefixes use lowercase a to z and are 1 to 100 characters long. There are up to 50,000 calls in total, in any mix.

Clarifying questions

• Which characters can appear? Lowercase letters only. (That allows an array of 26 slots; a dict works either way.)
• Can the same word be inserted twice? Yes, and it must behave as one word.
• Can a prefix or word be empty? No — at least one character.
• Is memory or speed the priority? Speed. There may be tens of thousands of prefix queries.

Approach 1: the simple way

Keep the words in a set. search is a set lookup. starts_with checks every stored word.

1class SetPrefixStore:2    """Brute force: a set of words; prefix queries scan every word."""34    def __init__(self) -> None:5        self.words: set[str] = set()67    def insert(self, word: str) -> None:8        self.words.add(word)910    def search(self, word: str) -> bool:11        return word in self.words1213    def starts_with(self, prefix: str) -> bool:14        return any(w.startswith(prefix) for w in self.words)

insert and search cost O(L). starts_with costs O(n × P) for n stored words and a prefix of length P, because it may compare the prefix with every word.

Put numbers on it. Suppose 25,000 inserts, then 25,000 prefix queries of up to 100 characters. The worst case is 25,000 × 25,000 × 100 ≈ 6 × 10^10 character comparisons. Even when most comparisons fail at the first letter, each query still visits all 25,000 words: 6 × 10^8 word visits. On a real 100,000-word list, one failed prefix query took about 5 milliseconds in Python. That is the price of reading every word to answer a question about a few letters.

The key insight

A prefix query only cares about words that share the first letter. Among those, it only cares about the ones that share the second letter, and so on. So group the words by first letter, then each group by second letter, then by third — nested groups, one level per character.

Those nested groups are a trie. Each group is a node, and each node stands for one prefix. To answer starts_with("ca"), you do not look at words at all: you step into the c group, then into its a group. If both exist, some word starts with ca. The number of stored words has vanished from the cost.

The one thing grouping loses is where each word ends. car and cart both pass through the car group. So each node carries a flag: "a word ends here".

Approach 2: a trie

1class TrieNode:2    """One prefix: its children by letter, and whether a word ends here."""34    def __init__(self) -> None:5        self.children: dict[str, TrieNode] = {}6        self.is_word = False789class Trie:10    """Prefix tree: every operation costs the length of its argument."""1112    def __init__(self) -> None:13        self.root = TrieNode()1415    def insert(self, word: str) -> None:16        node = self.root17        for ch in word:18            if ch not in node.children:19                node.children[ch] = TrieNode()20            node = node.children[ch]21        node.is_word = True2223    def _walk(self, prefix: str) -> TrieNode | None:24        node = self.root25        for ch in prefix:26            node = node.children.get(ch)27            if node is None:28                return None29        return node3031    def search(self, word: str) -> bool:32        node = self._walk(word)33        return node is not None and node.is_word3435    def starts_with(self, prefix: str) -> bool:36        return self._walk(prefix) is not None

Step by step:

1. insert walks from the root. For each character it goes down the matching edge, creating the node first if it is missing. After the last character it sets the flag.
2. _walk follows a string from the root and returns the node it ends on, or None as soon as an edge is missing.
3. search needs the path to exist and the flag to be set.
4. starts_with needs only the path.

Dry run on the example. The trie starts as a lone root.

The finished trie has 9 nodes including the root, against 13 characters in the four words.

Complexity. insert, search and starts_with each cost O(L) time for a string of length L, with no dependence on how many words are stored. Space is O(S) for S total inserted characters in the worst case, and less when prefixes are shared. On the same 100,000-word list, the failed prefix query that took 5 milliseconds with the set took about 1 microsecond with the trie.

Approach 3: fixed-array nodes

When the alphabet is known and small — here, a to z — each node can hold a list of 26 slots instead of a dict. The child for ch lives at index ord(ch) - ord("a").

1class ArrayTrieNode:2    """Node with 26 fixed slots, for lowercase a-z only."""34    __slots__ = ("children", "is_word")56    def __init__(self) -> None:7        self.children: list[ArrayTrieNode | None] = [None] * 268        self.is_word = False91011class ArrayTrie:12    """The same trie, indexing a 26-slot list instead of a dict."""1314    def __init__(self) -> None:15        self.root = ArrayTrieNode()1617    def insert(self, word: str) -> None:18        node = self.root19        for ch in word:20            i = ord(ch) - ord("a")21            if node.children[i] is None:22                node.children[i] = ArrayTrieNode()23            node = node.children[i]24        node.is_word = True

_walk, search and starts_with change the same way: index the list instead of calling .get. The complexity is unchanged. What changes is the constant factor. In C++ or Java the array is the classic choice: no hashing, and a node is one small block of memory. In Python the picture flips. On the 100,000-word list, the array version took about 124 MB against 98 MB for dict nodes, because most nodes have only one child and still carry 26 slots. Say this trade-off out loud; the wrong answer is having no reason for your choice.

Edge cases

• A stored word that is a prefix of another (car and cart). Node r is both an end and a pass-through. The flag handles it; "has no children" would not.
• Searching for a prefix of a stored word (search("ca")). Must be False. This is the test that catches a missing or misplaced flag.
• The same word inserted twice. The second insert walks the existing path and sets a flag that is already set. Nothing breaks.
• A query longer than any word (starts_with("cartwheel")). _walk returns None at the first missing edge.

Follow-ups

• "How many words start with this prefix?" Add a count to each node and increase it along the path when a new word is inserted. The query is a walk plus one read: O(P).
• "Return every word with this prefix." Walk to the prefix node, then DFS below it, collecting words at flagged nodes. Cost: O(P + size of that subtree).
• "Support delete." Clear the flag, then remove nodes on the way back up that have no children and no flag:

1    def delete(self, word: str) -> None:2        """Remove word if present, pruning nodes that no longer lead anywhere."""34        def prune(node: TrieNode, depth: int) -> bool:5            if depth == len(word):6                node.is_word = False7            else:8                child = node.children.get(word[depth])9                if child is None:10                    return False                 # word not stored; change nothing11                if prune(child, depth + 1):12                    del node.children[word[depth]]13            return not node.is_word and not node.children1415        prune(self.root, 0)

Deleting cart removes only node t: node r still ends car, so it stays.