Tries for Prefix Search

Trees keyed by characters

A trie (prefix tree) stores strings character by character: each node represents a prefix, children are keyed by the next character, and a flag marks nodes where a complete word ends. Inserting or looking up a word of length m costs O(m), independent of how many words are stored, and the defining strength is prefix queries: finding all words that start with "pun" means walking three nodes and then collecting the subtree. Tries power autocomplete, spell checkers, IP routing tables (longest-prefix match on bits), word games and dictionary compression. In Python, a node is naturally a dict of children plus an end-of-word marker, or nested dicts. Tries use more memory than a set of strings because of the many small nodes; radix trees (compressed tries) merge chains of single-child nodes to save space. Simple alternatives for prefix search over a static word list are a sorted list with bisect (find the first word ≥ prefix and scan while words start with it), which is compact and fast in Python.

An autocomplete trie with counts

Insert words with frequencies and return the most popular completions.

from __future__ import annotations
import heapq
from dataclasses import dataclass, field

@dataclass
class TrieNode:
    children: dict[str, TrieNode] = field(default_factory=dict)
    count: int = 0                         # > 0 means a complete word ends here

class Autocomplete:
    def __init__(self) -> None:
        self.root = TrieNode()

    def add(self, word: str, times: int = 1) -> None:   # O(len(word))
        node = self.root
        for ch in word.lower():
            node = node.children.setdefault(ch, TrieNode())
        node.count += times

    def suggest(self, prefix: str, k: int = 3) -> list[str]:
        node = self.root
        for ch in prefix.lower():
            node = node.children.get(ch)
            if node is None:
                return []
        found = []                                       # collect words under the prefix
        stack = [(node, prefix.lower())]
        while stack:
            current, text = stack.pop()
            if current.count:
                found.append((current.count, text))
            for ch, child in current.children.items():
                stack.append((child, text + ch))
        return [w for _, w in heapq.nlargest(k, found)]

ac = Autocomplete()
for city, searches in [("pune", 90), ("punjab", 40), ("puducherry", 25), ("patna", 30), ("pun", 5)]:
    ac.add(city, searches)
print(ac.suggest("pu"))       # ['pune', 'punjab', 'puducherry']
print(ac.suggest("pat"))      # ['patna']
print(ac.suggest("x"))        # []

Consider a sorted list first

For a few thousand static words, bisect on a sorted list gives prefix search with far less memory and code. Reach for a trie when the vocabulary is large, changes often, or you need per-prefix data such as counts.

Quick check: What is the cost of looking up a word of length m in a trie?

• O(n) in the number of words
• O(log n)
• O(m), independent of the number of stored words
• O(n·m)