Implement a Trie: Insert, Search and startsWith Explained

"Implement a trie with insert, search and startsWith." It is one of the few interview questions that asks you to build a data structure rather than use one, and the whole difficulty sits in a single boolean. Get the end-of-word flag right and the rest is a loop over characters.

The problem it solves

Store a set of words so that three questions are fast:

• insert(word) adds a word.
• search(word) says whether that exact word was inserted.
• starts_with(prefix) says whether any inserted word begins with prefix.

A hash set answers search in expected O(L) time for a word of length L, since it has to hash the string. It has no good answer for starts_with: without extra structure, it must check every stored word. A trie answers both by walking the same path.

The intuition

A trie spells words out instead of storing them whole. Each edge carries one character, so the path from the root down to any node spells a prefix. Words that share a prefix share those nodes: car, cart and cat all go through c and then a, and only split below that.

Two consequences follow:

• Every question is a walk. To search or check a prefix, start at the root and follow one edge per character. If an edge is missing, the answer is no, and you stop early.
• The path is not enough for search. After inserting cart, the path c, a, r exists, but car was never inserted. So each node carries a flag, is_end, set only where an inserted word stops. search needs the path and the flag; starts_with needs only the path.

Watch it run

The animation starts from a lone root and inserts car, creating one node per character and flagging r. Inserting cart walks the three existing nodes and builds only t. Inserting cat shares c and a and grows a second branch. The last two frames run the lesson's searches: search("car") ends on a flagged node, so it is a word; search("ca") finds the path but no flag, so it is only a prefix.

The code

A class-based trie. Each node holds a dictionary of children and the end-of-word flag; _walk is shared by both lookups:

class TrieNode:
    __slots__ = ("children", "is_end")
    def __init__(self):
        self.children = {}
        self.is_end = False

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

    def insert(self, word):
        node = self.root
        for ch in word:
            if ch not in node.children:
                node.children[ch] = TrieNode()
            node = node.children[ch]
        node.is_end = True

    def _walk(self, s):
        node = self.root
        for ch in s:
            node = node.children.get(ch)
            if node is None:
                return None
        return node

    def search(self, word):
        node = self._walk(word)
        return node is not None and node.is_end

    def starts_with(self, prefix):
        node = self._walk(prefix)
        # the root exists even in an empty trie, so "" needs a real word below
        return node is not None and (node.is_end or bool(node.children))

t = Trie()
for w in ["car", "cart", "cat"]:
    t.insert(w)
print(t.search("car"), t.search("ca"), t.starts_with("ca"), t.starts_with("co"))
# True False True False

Autocomplete is the same walk followed by a depth-first search below the prefix node. Visiting children in sorted order returns the words alphabetically:

def words_with_prefix(trie, prefix):
    node, out = trie._walk(prefix), []
    def dfs(n, path):
        if n.is_end:
            out.append(path)
        for ch in sorted(n.children):
            dfs(n.children[ch], path + ch)
    if node:
        dfs(node, prefix)
    return out

print(words_with_prefix(t, "ca"))        # ['car', 'cart', 'cat']

All three operations against a plain Python set as the reference, on 500 random word sets over a two-letter alphabet, 20 queries each. A small alphabet makes shared prefixes and prefix-only paths common, and the empty string shows up as both a word and a query:

import random

random.seed(8)
ok = True
for _ in range(500):
    words = {"".join(random.choice("ab") for _ in range(random.randint(0, 5)))
             for _ in range(random.randint(0, 12))}
    trie = Trie()
    for w in words:
        trie.insert(w)
    for _ in range(20):
        q = "".join(random.choice("ab") for _ in range(random.randint(0, 6)))
        ok &= trie.search(q) == (q in words)
        ok &= trie.starts_with(q) == any(w.startswith(q) for w in words)
        ok &= words_with_prefix(trie, q) == sorted(w for w in words if w.startswith(q))
print(ok)                                # True

A common shorter starts_with returns just _walk(prefix) is not None. The random test fails on it: an empty trie still has a root, so starts_with("") says yes with no words stored. That is the kind of case a hand-picked example never reaches.

The complexity

• insert, search, starts_with each take O(L) time for a string of length L: one dictionary step per character. The number of stored words never appears in the cost.
• Space is at most one node per inserted character, O(total characters), and shared prefixes make it smaller. Each node, though, is a whole object with its own dictionary, so in Python a trie usually takes more memory than a set of the same strings.
• Autocomplete costs O(P) to reach the prefix node plus time proportional to the size of the subtree below it. Sorting the children at each node adds a small factor; with a fixed alphabet, an array of 26 child slots gives the order for free.

Where it goes wrong

• No end-of-word flag. Without it, search("car") after inserting only cart returns true. The flag is the difference between a word and a prefix.
• Setting the flag on the wrong node. It belongs on the node reached after the last character, not on its parent.
• Treating the root as a word. The root spells the empty string. It is a word only if "" was inserted, and it is a prefix of something only if the trie is non-empty.
• Deleting by removing nodes blindly. Deleting car must only clear its flag while cart still needs the path; nodes can be pruned only when they have no children and no flag.
• Assuming the trie is always the better tool. For exact-match lookups alone, a hash set is simpler and usually lighter. The trie earns its memory when prefixes matter.

How to say it in an interview

"A trie stores words one character per edge, so the path from the root spells a prefix and shared prefixes share nodes. Each node has a child map and an end-of-word flag. Insert walks the word, creating missing children, and sets the flag on the last node. Search walks the word and needs both the path and the flag; startsWith needs only the path. All three are O(L) in the length of the string, independent of how many words are stored. Memory is at most one node per inserted character, less when prefixes are shared."

When to reach for a trie over a set is weighed in Trie vs hash set for prefixes, and the autocomplete walk is its own lesson, prefix search and autocomplete.