algo-trie-autocomplete

Answer

The core data structure is a prefix tree (Trie) where each node stores children, a word-end flag, and a frequency counter. I provide two strategies optimized for different large-vocabulary workloads:

Strategy A – DFS + Sort (write-optimized)

Best when inserts dominate. No extra bookkeeping. Autocomplete collects all matching words via DFS, sorts by freq, and returns top K.

from typing import List, Optional, Tuple

class TrieNode:
    __slots__ = ('children', 'freq', 'is_end')
    def __init__(self):
        self.children: dict = {}
        self.freq: int = 0
        self.is_end: bool = False


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

    def insert(self, word: str, freq: int = 1) -> None:
        node = self.root
        for ch in word:
            if ch not in node.children:
                node.children[ch] = TrieNode()
            node = node.children[ch]
        node.is_end = True
        node.freq += freq

    def search(self, word: str) -> bool:
        node = self._find(word)
        return node is not None and node.is_end

    def starts_with(self, prefix: str) -> bool:
        return self._find(prefix) is not None

    def autocomplete(self, prefix: str, k: int = 10) -> List[str]:
        node = self._find(prefix)
        if node is None:
            return []
        results: List[Tuple[str, int]] = []
        self._collect(node, prefix, results)
        results.sort(key=lambda x: (-x[1], x[0]))
        return [w for w, _ in results[:k]]

    def _find(self, prefix: str) -> Optional[TrieNode]:
        node = self.root
        for ch in prefix:
            if ch not in node.children:
                return None
            node = node.children[ch]
        return node

    def _collect(self, node: TrieNode, prefix: str, results: list) -> None:
        if node.is_end:
            results.append((prefix, node.freq))
        for ch, child in node.children.items():
            self._collect(child, prefix + ch, results)

Why it works for large vocabularies:
- O(L) insert/search (L = word length).
- O(N log N) autocomplete in worst case (N = matching words), but for real vocabularies with well-distributed prefixes N is small.
- No memory overhead beyond the trie itself.

Strategy B – Pre-computed Top-K at Each Node (read-optimized)

When autocomplete queries vastly outnumber inserts. Each node caches its top K completions; every insert propagates an update bottom-up so queries are instant.

from typing import List, Optional, Tuple

class TrieNode:
    __slots__ = ('children', 'freq', 'is_end', 'top_k')
    def __init__(self):
        self.children: dict = {}
        self.freq: int = 0
        self.is_end: bool = False
        self.top_k: List[Tuple[str, int]] = []


class Trie:
    def __init__(self, k: int = 10):
        self.root = TrieNode()
        self.k = k

    def insert(self, word: str, freq: int = 1) -> None:
        node = self.root
        path = [node]
        for ch in word:
            if ch not in node.children:
                node.children[ch] = TrieNode()
            node = node.children[ch]
            path.append(node)
        node.is_end = True
        node.freq += freq
        for n in reversed(path):
            self._refresh_top_k(n)

    def _refresh_top_k(self, node: TrieNode) -> None:
        candidates: List[Tuple[str, int]] = []
        if node.is_end:
            candidates.append(("", node.freq))
        for ch, child in node.children.items():
            for suffix, f in child.top_k:
                candidates.append((ch + suffix, f))
        candidates.sort(key=lambda x: (-x[1], x[0]))
        node.top_k = candidates[: self.k]

    def search(self, word: str) -> bool:
        node = self._find(word)
        return node is not None and node.is_end

    def starts_with(self, prefix: str) -> bool:
        return self._find(prefix) is not None

    def autocomplete(self, prefix: str, k: Optional[int] = None) -> List[str]:
        k = k or self.k
        node = self._find(prefix)
        if node is None:
            return []
        if len(node.top_k) >= k:
            return [prefix + s for s, _ in node.top_k[:k]]
        results: List[Tuple[str, int]] = []
        self._collect(node, prefix, results)
        results.sort(key=lambda x: (-x[1], x[0]))
        return [w for w, _ in results[:k]]

    def _collect(self, node, prefix, results):
        if node.is_end:
            results.append((prefix, node.freq))
        for ch, child in node.children.items():
            self._collect(child, prefix + ch, results)

    def _find(self, prefix: str) -> Optional[TrieNode]:
        node = self.root
        for ch in prefix:
            if ch not in node.children:
                return None
            node = node.children[ch]
        return node

Why it works for large vocabularies:
- autocomplete() is O(K) — constant time regardless of vocabulary size.
- Trade-off: insert is O(L × B × K) where B is average branching factor, K is top-k limit.
- Use when queries per insert ratio > ~100:1.

Choice Recommendation