Autocomplete With a Trie: Prefix Search and Top-K Suggestions

A spell checker asks "is this a word?". Autocomplete asks something harder: "which stored words start with what has been typed so far, and which few of them should I show?". A trie answers the first half of that in time proportional to the prefix, however large the dictionary. The second half, picking a handful of suggestions without reading thousands of words, is where most implementations are slower than they need to be.

The problem it solves

Given a dictionary, and for each keystroke a prefix, return up to k suggestions. Two orderings are common:

• Alphabetical: the first k matching words in sorted order. This is the standard coding question, usually with k = 3 and one list per typed character.
• By popularity: the k matching words with the highest counts, ties broken alphabetically. This is what a search box does, and it is where the design question starts.

A hash set cannot help with either: hashing scatters words that share a prefix, so every keystroke means scanning the whole dictionary. A sorted list with a binary-searched prefix range works for a static dictionary, as trie vs hash set shows. A trie makes each prefix a node, which pays off once you want to cache per prefix.

The intuition

Every node in a trie is a prefix: the path from the root spells it. So autocomplete is two steps:

1. Walk the prefix. One hop per typed character. The cost is O(P) for a prefix of length P, and the number of stored words never enters it. If a character has no child, no stored word starts that way and the answer is empty.
2. Read the subtree. Everything below the prefix node shares the prefix, and nothing outside it does.

Step two is the trap. A long prefix has a tiny subtree; a one-letter prefix can hold a large slice of the dictionary, and collecting all of it per keystroke defeats the point. Two fixes:

• Alphabetical order: traverse children in sorted order and stop after k words. The first word found in a sorted depth-first walk is the smallest, so there is no need to collect the rest.
• Popularity order: there is no early stop, because the most popular word could be deep in any branch. So precompute: every node keeps a small cached list of the best k words in its subtree, updated on insert. A query becomes a walk plus a read of that list.

Typing adds one more saving. The node for "ca" is the parent of the node for "cat", so keeping the current node between keystrokes makes each new character a single hop instead of a fresh walk from the root.

Watch it run

The animation stores car, cart and cat. With nothing typed, every one of them is still a candidate. The user types c: one hop down from the root. Then a: two characters typed, two hops walked, and the dictionary size never entered it. That node is the prefix ca, and everything in the subtree below it starts with those two letters. Collecting the end-of-word flags finds car at r, then cart further down the same branch, then cat on the other branch out of a. Three suggestions for two hops, and nothing outside the ca subtree was ever read. Finally the user types cx: the node for c has no x child, so the walk stops after one hop with no suggestions.

The code

One trie that answers all three questions: every word under a prefix, the first k alphabetically with an early stop, and the k most popular from a per-node cache:

class Node:
    __slots__ = ("kids", "word", "top")
    def __init__(self):
        self.kids = {}        # char -> Node
        self.word = None      # set when a stored word ends here
        self.top = []         # best (-count, word) pairs in this subtree, cached

class Autocomplete:
    def __init__(self, counts, k=3):
        self.root, self.k = Node(), k
        for word, count in counts.items():
            self._insert(word, count)

    def _insert(self, word, count):
        node, path = self.root, [self.root]
        for ch in word:
            node = node.kids.setdefault(ch, Node())
            path.append(node)
        node.word = word
        for n in path:                      # every prefix of word may now rank it
            n.top = sorted(n.top + [(-count, word)])[: self.k]

    def walk(self, prefix):                 # O(len(prefix)); None if the path breaks
        node = self.root
        for ch in prefix:
            node = node.kids.get(ch)
            if node is None:
                return None
        return node

    def first_k(self, node, k):             # sorted depth-first walk, stops early
        out, stack = [], [node] if node else []
        while stack and len(out) < k:
            node = stack.pop()
            if node.word is not None:
                out.append(node.word)
            stack.extend(node.kids[c] for c in sorted(node.kids, reverse=True))
        return out

    def all_under(self, prefix):            # the lesson's walk-then-collect
        return self.first_k(self.walk(prefix), float("inf"))

    def most_popular(self, prefix):         # O(len(prefix) + k): read the cache
        node = self.walk(prefix)
        return [w for _, w in node.top] if node else []

counts = {"car": 50, "cart": 20, "cat": 90, "care": 70, "cast": 5, "dog": 40}
ac = Autocomplete(counts)
print(ac.all_under("ca"))       # ['car', 'care', 'cart', 'cast', 'cat']
print(ac.most_popular("ca"))    # ['cat', 'care', 'car']
print(ac.most_popular("car"))   # ['care', 'car', 'cart']
print(ac.most_popular("cx"))    # []

The children are pushed in reverse sorted order so the stack pops them smallest first, which makes the walk alphabetical. Per keystroke, the previous node is kept and each character costs one hop:

def suggest_as_typed(ac, typed, k=3):
    """One hop per keystroke: the node for the previous prefix is kept."""
    node, rows = ac.root, []
    for ch in typed:
        node = node.kids.get(ch) if node else None
        rows.append(ac.first_k(node, k))
    return rows

print(suggest_as_typed(ac, "cas"))
# [['car', 'care', 'cart'], ['car', 'care', 'cart'], ['cast']]

Checked on 300 seeded random dictionaries with 20 random prefixes each, including prefixes with letters no word uses, against a brute force that scans every word with startswith:

import random

random.seed(28)
ok = True
for _ in range(300):
    words = {"".join(random.choice("abc") for _ in range(random.randint(1, 6))): random.randint(1, 99)
             for _ in range(random.randint(1, 40))}
    ac = Autocomplete(words, k=3)
    for _ in range(20):
        p = "".join(random.choice("abcd") for _ in range(random.randint(0, 4)))
        hits = sorted(w for w in words if w.startswith(p))       # brute force: scan all
        ok &= ac.all_under(p) == hits
        ok &= ac.first_k(ac.walk(p), 3) == hits[:3]
        ok &= ac.most_popular(p) == sorted(hits, key=lambda w: (-words[w], w))[:3]
print(ok)                        # True

The complexity

With P the prefix length, L the word length and k the number of suggestions:

• Walk: O(P), independent of the dictionary size. The Big-O cheat sheet lists the trie's starts-with at O(L) for the same reason.
• All words under a prefix: O(P + size of the subtree).
• First k alphabetically: O(P) plus the nodes visited before the k-th word is found, times the cost of sorting each node's children (bounded by the alphabet).
• Top k by popularity, cached: O(P + k) per query. Insert pays O(L · k log k) to refresh the caches along the path.
• Space: one node per distinct prefix, plus k cached entries per node when caching.

Where it goes wrong

• Collecting the whole subtree for short prefixes. A one-letter prefix can match most of the dictionary; stop early or cache.
• Updating a count without removing the old cache entry. Re-inserting a word with a new count leaves the stale pair behind; replace the word's entry in each cache on its path.
• Walking from the root on every keystroke. Keep the node; typing one more character is one more hop.
• Forgetting the empty answer. A broken path means nothing matches; return an empty list rather than falling back to the root.

When it shows up in interviews

Two ways. As a coding question: build a trie, then for every typed prefix return up to three matching words in alphabetical order, where the early stop and the kept node are the points being tested. And as a design question, design search autocomplete, where the per-node top-k cache is the core idea and the follow-ups are how counts are updated offline and how the trie is sharded. The trie itself is built in implement a trie, and the same walk-and-prune idea drives word search II.

How to say it in an interview

"Every trie node is a prefix, so I walk the typed characters, one hop each, which is O(P) regardless of dictionary size. If a character is missing, nothing matches. The subtree below that node holds exactly the matching words. For alphabetical suggestions I do a depth-first walk with children in sorted order and stop after k words. For popularity I cache the top k words at every node when inserting, so a query is the walk plus reading k entries. Between keystrokes I keep the current node, so each new character is a single hop."