Approach:Use a Trie where each node stores:- children map- is_word flag- optionally word and score- a fixed-size min-heap (or sorted list) of top-K entries for that subtree (store (score, word) with score positive; min-heap keeps smallest at root so we evict lowest when >K). On insert, walk the trie and update each node's top-K structure.This lets suggest(prefix, k) reach the node for prefix in O(|prefix|) and return its top-K in O(k log k) (or O(k) if maintained sorted).Code (supports variable k per query by storing up to global_K at nodes — choose a reasonable max_K like 50):python
import heapq
from typing import Dict, List, Tuple
class TrieNode:
def __init__(self, max_k):
self.children: Dict[str, TrieNode] = {}
self.is_word = False
self.word = None
self.score = 0
self.max_k = max_k
# min-heap of (score, word). We keep size <= max_k
self.topk: List[Tuple[int, str]] = []
def push_topk(self, score, word):
entry = (score, word)
if len(self.topk) < self.max_k:
heapq.heappush(self.topk, entry)
else:
# replace smallest if new score better or tie but lexicographically smaller not required
if entry > self.topk[0]:
heapq.heapreplace(self.topk, entry)
class AutocompleteTrie:
def __init__(self, max_k=50):
self.root = TrieNode(max_k)
self.max_k = max_k
def insert(self, word: str, score: int):
node = self.root
node.push_topk(score, word)
for ch in word:
if ch not in node.children:
node.children[ch] = TrieNode(self.max_k)
node = node.children[ch]
node.push_topk(score, word)
node.is_word = True
node.word = word
node.score = score
def suggest(self, prefix: str, k: int) -> List[Tuple[int, str]]:
node = self.root
for ch in prefix:
if ch not in node.children:
return []
node = node.children[ch]
# node.topk is a min-heap; to return descending scores, sort descending
items = sorted(node.topk, key=lambda x: (-x[0], x[1]))
return items[:k]
Key points:- Maintaining top-K at each node makes suggest O(|prefix| + k log k) to sort/format.- Insert cost: O(|word| * log K) because at each node we push/replace in heap.- Space: additional O(N * K) for storing top-K across nodes in worst case (N = total nodes). That can be large.Trade-offs:- Pros: Very fast suggestions, simple to implement, good for read-heavy systems.- Cons: High memory overhead; updates (insert/score changes) are more expensive since you must update all ancestor nodes' heaps. Deleting or decreasing score is tricky — may require re-computing top-K for affected nodes or using more complex data structures (balanced BSTs, skip lists, or storing references to entries).- Alternatives: Don't store top-K per node; instead, on suggest, traverse subtree and use a heap to collect top-K — saves memory but costs time (O(nodes_in_subtree + log k)). Or store inverted index of prefixes to candidates (more memory but faster retrieval). For dynamic score updates, consider maintaining a global mapping and storing pointers in node buckets with lazy re-heapifying.Edge cases:- Choose max_k >= maximum requested k. If query k > stored max_k, either increase stored K (rebuild) or perform a subtree search at query time.- Ties in score — break deterministically (lexicographic).- Large vocabulary: consider compressing trie (radix trie) or sharding by prefix.