26. Trie-Based Vocabulary Autocomplete

Problem

A research tool autocompletes tokenizer vocabulary entries as a researcher types a prefix while inspecting model outputs.

Given the full vocabulary and a prefix, build a trie over the vocabulary once, then return up to limit matching vocabulary entries that start with the prefix, sorted alphabetically. The walk to the prefix's node must be O(len(prefix)) — it must not degrade to scanning the entire vocabulary per query.

Source: src/26_trie_prefix_autocomplete.py

def autocomplete_suggestions(vocabulary: list[str], prefix: str, limit: int) -> list[str]:
    ...

Examples:
>>> autocomplete_suggestions(["cat", "car", "cart", "dog", "care"], "car", 3)
['car', 'care', 'cart']

>>> autocomplete_suggestions(["cat", "car"], "do", 5)
[]

Step-by-Step Approach

1. Recognize "repeated prefix lookups over a fixed vocabulary" as the signature trie use case — a trie turns each query's cost into O(len(prefix) + number of matches), independent of how large the vocabulary is.
2. Build the trie once: for each word, walk character by character from the root, creating a child node for any character not yet present, and mark the final node as a complete word.
3. To answer a query, walk from the root following the prefix's characters one at a time. If any character is missing from the trie at that point, no vocabulary entry starts with this prefix — return an empty list immediately.
4. Once at the prefix's node, every complete word reachable below it (including the node itself, if it's a complete word) is a match — collect them with a DFS.
5. Visit children in sorted character order during the DFS, and emit a node's word (if any) before descending into its children — a completed word at a node is always lexicographically no greater than anything formed by continuing deeper from that same node, so this ordering produces alphabetical output with no separate sort step.
6. Stop early — both during word collection and during the DFS descent — the moment `limit` results have been gathered, avoiding unnecessary work when the vocabulary has far more matches than are needed.

The key insight is that building the trie is a one-time O(total vocabulary length) cost, after which every prefix query only pays for the characters in the prefix plus the matches it returns — a linear scan over the vocabulary would instead pay O(vocabulary size) on every single query.

Reference solution

class _TrieNode:
    __slots__ = ("children", "is_word")

    def __init__(self):
        self.children: dict[str, "_TrieNode"] = {}
        self.is_word = False


def autocomplete_suggestions(vocabulary: list[str], prefix: str, limit: int) -> list[str]:
    root = _TrieNode()
    for word in vocabulary:
        node = root
        for ch in word:
            node = node.children.setdefault(ch, _TrieNode())
        node.is_word = True

    # walk to the prefix's node: O(len(prefix)), independent of vocabulary size
    node = root
    for ch in prefix:
        if ch not in node.children:
            return []
        node = node.children[ch]

    # DFS from here, always visiting children in sorted order and emitting a word the
    # instant we pass through a node marked is_word. Because a completed word at a node
    # is always lexicographically <= any word formed by continuing deeper from that same
    # node, this naturally yields alphabetical order with no extra sort needed.
    results: list[str] = []

    def dfs(n: _TrieNode, path: str) -> None:
        if len(results) >= limit:
            return
        if n.is_word:
            results.append(prefix + path)
        for ch in sorted(n.children):
            if len(results) >= limit:
                return
            dfs(n.children[ch], path + ch)

    dfs(node, "")
    return results[:limit]

Key Functions & Tricks

• node.children.setdefault(ch, _TrieNode()) — creates a child on first visit, reuses it on every later word that shares this prefix
• is_word flag — distinguishes "this exact string is a vocabulary entry" from "this is just an intermediate prefix of longer entries"
• Prefix walk with early return [] — the O(len(prefix)) step that makes the whole query fast regardless of vocabulary size
• DFS over sorted(n.children) — visiting children in character order makes the output alphabetical for free
• Emit-before-descend ordering — a shorter completed word always sorts before longer words that extend it, matching Python string comparison
• len(results) >= limit early-exit checks — stop collecting the moment enough matches are found

How to Recognize This Pattern