Trie in Python: Prefix Search and Wildcard Matching

Type “ap” into a search box and suggestions appear before you finish the word. Behind that behaviour sits a surprisingly small structure: a tree where every step down spells one more letter. A trie (pronounced “try”, also called a prefix tree) turns “does any word start with this?” from a scan over every word into a walk whose length equals the prefix, no matter how many words are stored.

This guide builds one, adds autocomplete and wildcard search, and then does something many tutorials skip: it measures a trie against simpler tools, so you know when it is actually the right answer. All code was run, with a 2,000-case random test against a plain set and a 16,000-query test of the wildcard search against Python’s re module.

What a trie actually stores

Each node holds a dictionary from the next letter to a child node, plus one boolean saying whether a word ends here. Words that share a beginning share a path. The figure shows the words app, apple, apt and bat.

Notice that “app” is both a complete word and a stepping stone to “apple”. That is why the boolean flag exists. The shape of the path alone cannot tell the difference.

The implementation

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

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

    def insert(self, word):
        node = self.root
        for ch in word:
            node = node.children.setdefault(ch, TrieNode())
        node.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):
        n = self._walk(word)
        return n is not None and n.end

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

The helper _walk is the design decision worth copying. Both search and starts_with are “walk the string, then ask a question about where you landed”, so write the walk once. After inserting the seven words apple, app, apply, apt, bat, bath and batman:

t.search('app')       # -> True
t.search('appl')      # -> False   (a prefix, but no word ends there)
t.starts_with('appl') # -> True
t.starts_with('cat')  # -> False

The bug that passes the first test

If you forget the end flag and make search simply return whether the walk succeeded, every inserted word still works, so basic tests pass. Only a prefix query exposes it. With just “apple” inserted, the broken version printed search('app') -> True. Always test a word that is a strict prefix of another.

Autocomplete: walk, then explore

To list suggestions, walk to the prefix node, then run a depth-first search below it, visiting children alphabetically so the output is predictable. Stop at a limit so a short prefix does not return the whole dictionary.

def autocomplete(self, prefix, limit=3):      # a method of Trie
    node = self._walk(prefix)
    out = []
    if node is None:
        return out

    def dfs(n, path):
        if len(out) >= limit:
            return
        if n.end:
            out.append(prefix + ''.join(path))
        for ch in sorted(n.children):
            path.append(ch)
            dfs(n.children[ch], path)
            path.pop()
            if len(out) >= limit:
                return

    dfs(node, [])
    return out

# t.autocomplete('ap')  -> ['app', 'apple', 'apply']
# t.autocomplete('bat') -> ['bat', 'bath', 'batman']
# t.autocomplete('z')   -> []

The path.append and path.pop pair is the same choose-and-undo rhythm used in the backtracking guide, and the recursion itself is the depth-first traversal from the tree traversal article.

Wildcard search: where tries earn their keep

Add and Search Words lets a query contain ., which matches any single letter. A set cannot do that without checking every word, but a trie can branch: on a normal letter, follow one child; on a dot, try all children.

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

    def addWord(self, w):
        n = self.root
        for ch in w:
            n = n.children.setdefault(ch, TrieNode())
        n.end = True

    def search(self, w):
        def dfs(i, node):
            if i == len(w):
                return node.end
            ch = w[i]
            if ch == '.':
                return any(dfs(i + 1, c) for c in node.children.values())
            c = node.children.get(ch)
            return c is not None and dfs(i + 1, c)
        return dfs(0, self.root)

With bad, dad and mad added, the six queries pad, bad, .ad, b.., … and …. returned False, True, True, True, True, False. A word of three letters never matches a four-letter pattern, because the recursion reaches the end of a path before the pattern ends. Against Python’s re.fullmatch on 16,000 random queries, every result was identical.

Do you even need a trie?

Here is the honest comparison. The test stored 99,784 random lowercase words and asked 300 times, “does any word start with this three-letter prefix?”, three different ways:

Absolute times depend on the machine, but the ratio does not: the trie is thousands of times faster than scanning. The surprise is the third row. For a static word list and a plain “any word with this prefix?” question, a sorted list with binary search is just as fast and far lighter. The same test built a trie with 406,739 nodes for 698,239 characters of text, and every node is a Python object plus a dictionary. The binary search guide covers the technique, with bisect summarised in the Python cheat sheet.

So choose a trie when you need what the sorted list cannot give you:

• Wildcards or pattern branching, as in Add and Search Words.
• Words arriving over time, where re-sorting after every insert would be wasteful.
• Data on the path, such as a count of words below each node or the most popular completion.
• Pruning a search, as in Word Search II, where the trie tells the grid search to stop the moment no word can continue.

Complexity

Insert, search and prefix check are O(L) for a word of length L, independent of how many words are stored. Memory is O(total characters) in the worst case, less when prefixes are shared. A fixed alphabet can use a 26-slot array per node for speed, but a dictionary is simpler and handles any characters. For the notation behind these statements, see Big O notation explained.

Problems to practise next

1. Implement Trie (Prefix Tree), the exact class above.
2. Add and Search Words Data Structure, the wildcard version.
3. Replace Words, where each word is replaced by its shortest dictionary root.
4. Search Suggestions System, which pairs a sorted list with bisect.
5. Word Search II, a trie combined with grid backtracking.

Want a related structure for grouping rather than prefixes? See Union-Find, and for a hash-based alternative built from nothing, read how a hash table works inside.