Trie in Python: how a prefix tree works and when a problem needs one

Short answer. A trie is a tree. Every path down from the root spells a prefix. Each node holds its children by letter and a flag that says "a word ends here". Insert and lookup cost O(L) for a word of length L. The number of stored words does not matter.

What is a trie?

Take the words "car", "cat" and "cart". A trie stores them once, letter by letter, sharing what they share.

• The root has one child, c.
• c has one child, a.
• a has two children, r and t. The t is marked as an end.
• r is marked as an end. It has one child, t, also an end.

Six nodes for three words. The prefix "ca" exists once, not three times.

The other name is prefix tree. Same thing.

How do you write a trie in Python?

A dictionary per node keeps the code short. No fixed alphabet, no arrays of 26.

class Trie:
    def __init__(self):
        self.root = {}

    def insert(self, word):
        node = self.root
        for ch in word:
            node = node.setdefault(ch, {})
        node["$"] = True  # a word ends here

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

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

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

Three details carry the code.

• The end marker is a key that no letter can be. "$" works; so does None.
• search and starts_with share one walk. Only the last check differs.
• Every method is O(L). The number of stored words never appears.

When does a problem want a trie?

The wording again.

• "Autocomplete", "suggest words as the user types".
• "Does any word start with this prefix" asked many times.
• "Search a board for all words from a list".
• "Longest common prefix" over many strings.
• "Replace words with their shortest root".
• Bits instead of letters: "maximum XOR of two numbers".

The common thread: many strings and many prefix questions. A hash set answers only whole words.

Signals, near misses and code templates are on the trie pattern page.

How is a trie different from a hash set?

Both answer "is this word here". They part ways on prefixes.

No prefix questions in the problem? Use the set. A trie you do not need is a hundred lines of nothing.

How do you solve Word Search II with a trie?

The problem. A grid of letters and a list of words. Return every word that can be traced through adjacent cells. Each cell is used once per word.

Searching the grid once per word is too slow. Put the words in a trie and search the grid once.

Run a DFS from every cell. Carry the trie node along. A cell whose letter is not a child of the current node is a dead end. The search stops there.

def find_words(board, words):
    trie = Trie()
    for w in words:
        trie.insert(w)

    found = set()
    rows, cols = len(board), len(board[0])

    def dfs(r, c, node, path):
        ch = board[r][c]
        if ch not in node:
            return
        nxt = node[ch]
        path += ch
        if "$" in nxt:
            found.add(path)
        board[r][c] = "#"
        for dr, dc in ((1, 0), (-1, 0), (0, 1), (0, -1)):
            nr, nc = r + dr, c + dc
            if 0 <= nr < rows and 0 <= nc < cols and board[nr][nc] != "#":
                dfs(nr, nc, nxt, path)
        board[r][c] = ch

    for r in range(rows):
        for c in range(cols):
            dfs(r, c, trie.root, "")
    return list(found)

The "#" marks a cell as used on the current path. It is put back on the way out.

The pruning is the whole point. Without the trie, every path is followed to its full length. With it, most paths die on the second letter.

How do you explain a trie in an interview?

Three sentences before the code.

• "There are many prefix questions, so a trie. A set only answers whole words."
• "Each node maps a letter to a child and marks where a word ends."
• "Insert and lookup are O(L). The word count does not matter."

Then the class. Then the three-word example above. Say what the end marker is for. That is the part people skip.

What goes wrong with a trie?

• No end marker. "car" and "cart" become one word, and search("car") is wrong.
• End marker as a letter key. A word containing that letter breaks the trie.
• Rebuilding the trie per query. Build once, then ask.
• A trie for a problem with no prefixes. Slower to write and no faster to run.

Recognising the pattern from the wording is the harder half. That is what LeetCode patterns is about.

FAQ

What is the time complexity of a trie?

O(L) for insert, search and prefix lookup. L is the word's length. Building it is O(total letters).

How much memory does a trie use?

One node per distinct prefix. Many words with shared prefixes are cheap. Random strings are expensive.

Trie or hash map for autocomplete?

Trie. Autocomplete is a prefix question. A hash map cannot answer it without scanning.

Can a trie store numbers?

Yes, bit by bit. A binary trie over 32 bits solves "maximum XOR of two numbers" in O(32 n).