The machine runs a typeahead. The corpus is a set of sentences, each with a times-seen count. Every input(c) appends to the current prefix, and the machine must return the top 3 sentences by (count, then lexicographic order) that start with that prefix. input('#') terminates the sentence: the machine increments its count and resets the prefix.
A trie is the core structure: nodes keyed by character, each carrying the best 3 sentences in its subtree. The trick is maintaining this top-3 cache incrementally. When a count changes, the machine walks the trie along the sentence's characters, updating cached lists at every node on that path — O(L) per update for sentence length L. Each node's list stays sorted by (count desc, sentence asc), and merging a candidate into a 3-element list is O(3).
The hot-score update is the trap. When the user types a complete sentence ending in '#', that sentence's count must rise even though its characters were never re-inserted. The machine stores count in the trie's terminal node and re-propagates up the path — O(L) per '#'.
Lookup: input(c) follows the prefix down the trie — O(L) per keystroke — and returns the node's cached list. Memory is O(total characters × 3). Edge cases: input('#') first (ignore the empty sentence); equal counts sort lexicographically, so the comparator is (count, sentence); an unknown prefix returns [].
class AutocompleteSystem:
def __init__(self, sentences, times):
self.trie = {}
self.counts = defaultdict(int)
self.prefix = ""
for s, t in zip(sentences, times):
self._insert(s, t)
def _insert(self, s, times):
self.counts[s] += times
node = self.trie
for ch in s:
node = node.setdefault(ch, {})
node['*'] = self._top3(node.get('*', []), s, self.counts[s])
node['$'] = self.counts[s]
def _top3(self, top3, s, c):
top3.append((c, s))
top3 = [x for x in top3 if x[1] != s or x[0] == c]
return sorted(top3, key=lambda x: (-x[0], x[1]))[:3]
def input(self, c):
if c == '#':
self._insert(self.prefix, 1)
self.prefix = ""
return []
self.prefix += c
node = self.trie
for ch in self.prefix:
if ch not in node:
return []
node = node[ch]
return [s for _, s in node.get('*', [])]