Implement Trie (Prefix Tree)

Medium·🌿 Tries·Open on LeetCode ↗

Build a data structure that stores words and can answer two questions fast: is this exact word stored, and is any stored word starting with this prefix.

insert "car", "cat", "cod"step 1/29

An empty trie is one root node holding nothing. Words become paths down from it.

1class Trie:
2    def __init__(self):
3        self.children = {}
4        self.is_word = False
5 
6    def insert(self, word):
7        node = self
8        for ch in word:
9            if ch not in node.children:
10                node.children[ch] = Trie()
11            node = node.children[ch]
12        node.is_word = True
13 
14    def _walk(self, prefix):
15        node = self
16        for ch in prefix:
17            if ch not in node.children:
18                return None
19            node = node.children[ch]
20        return node
21 
22    def search(self, word):
23        node = self._walk(word)
24        return node is not None and node.is_word
25 
26    def startsWith(self, prefix):
27        return self._walk(prefix) is not None

Read the 29 steps as text

1An empty trie is one root node holding nothing. Words become paths down from it.
2Insert "car" — start at the root.
3No "c" child yet. Create it and step in.
4No "a" child yet. Create it and step in.
5No "r" child yet. Create it and step in.
6End of "car" — flag this node is_word. Note the node itself doesn't spell the word; the path does.
7Insert "cat" — start at the root.
8"c" already exists here — walk into the node "c" instead of making a new one.
9"a" already exists here — walk into the node "ca" instead of making a new one.
10No "t" child yet. Create it and step in.
11End of "cat" — flag this node is_word. Note the node itself doesn't spell the word; the path does.
12Insert "cod" — start at the root.
13"c" already exists here — walk into the node "c" instead of making a new one.
14No "o" child yet. Create it and step in.
15No "d" child yet. Create it and step in.
16End of "cod" — flag this node is_word. Note the node itself doesn't spell the word; the path does.
17Walk "cat" from the root.
18"c" exists. Step into it.
19"a" exists. Step into it.
20"t" exists. Step into it.
21Landed on a node, and it is flagged is_word. search("cat") is true.
22Walk "ca" from the root.
23"c" exists. Step into it.
24"a" exists. Step into it.
25The walk succeeded, but this node has no is_word flag — "ca" is a prefix, not a stored word. search("ca") is false.
26Walk "ca" from the root.
27"c" exists. Step into it.
28"a" exists. Step into it.
29Same walk, same node — but startsWith never looks at is_word, so it is true. That flag is the whole difference.

The idea

A hash set answers 'is this word stored' in O(1) and is the right answer if that's the only question. It is useless for the second one: to find whether anything starts with 'ca' you'd have to scan every key. Prefixes are the reason a trie exists.

The idea is to stop storing words and start storing letters. Each node owns a map from a character to a child node, so a word is a path down from the root rather than an entry in a table. Words with a shared prefix share the path that spells it — 'car' and 'cat' walk the same two nodes before splitting.

That's what makes prefix queries cheap. 'Does anything start with ca' is just 'can I walk c, then a' — O(length of prefix), completely independent of how many words are stored.

The subtle part is that a node cannot tell you whether a word ends there. The path c-a exists after inserting 'car', but 'ca' is not a stored word. So each node carries an explicit is_word flag, and that flag is the *entire* difference between search and startsWith. Both walk identically; only the last line differs.

The approach

1Give every node a map from character to child node, plus a boolean is_word.
2To insert: walk the word one character at a time, creating a child whenever the character is missing, then mark the final node is_word.
3To walk a prefix: follow the characters, returning null the moment one is missing.
4search returns true when the walk lands somewhere *and* that node is flagged is_word.
5startsWith returns true when the walk lands anywhere at all.

Complexity

Time

O(L)

Space

O(total characters)

Every operation costs the length of the word, L, with no dependence on how many words are stored. Space is bounded by the total characters inserted, and is less in practice because shared prefixes are stored once.

Code

class Trie:
    def __init__(self):
        self.children = {}
        self.is_word = False

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

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

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

    def startsWith(self, prefix):
        return self._walk(prefix) is not None

What goes wrong

•Making search and startsWith the same function. They differ by the is_word check, and conflating them makes search return true for every prefix of every stored word.
•Forgetting to set is_word on the final node of an insert, which makes search return false for words you definitely inserted.
•Using a fixed 26-slot array without saying why. It's the faster choice for lowercase-only input and the standard answer, but it costs 26 pointers per node even for a sparse trie — a hash map is the better default when the alphabet is unknown.
•Deleting by unlinking a leaf's parent chain without checking is_word on the way up. Removing 'car' must not remove 'cat', which shares its first two nodes. Deletion isn't asked for here, but it's the natural follow-up question.