1"""Chapter 34 (case study: search autocomplete / typeahead) deep-dive code.
21. Trie with top-K stored at every node vs a precomputed prefix -> top-K table.
32. Building top-K per prefix from a query log, and the memory it takes.
43. Blending long-term popularity with trending queries (exponential decay).
54. Caching: how much traffic short prefixes carry.
6Run: python3 ch34_autocomplete.py"""
7import heapq
8import math
9import random
10from collections import Counter, defaultdict
11
12K = 5
13
14
15# ---- 1a. trie with top-K at each node ----------------------------------------------
16
17class Trie:
18 def __init__(self):
19 self.root = {"kids": {}, "top": []}
20 self.nodes = 1
21
22 def build(self, counts):
23 for q, c in counts.items():
24 node = self.root
25 for ch in q:
26 if ch not in node["kids"]:
27 node["kids"][ch] = {"kids": {}, "top": []}
28 self.nodes += 1
29 node = node["kids"][ch]
30 # fill top-K bottom-up: each node keeps the best K of its subtree
31 self._fill(self.root, "", counts)
32
33 def _fill(self, node, prefix, counts):
34 cands = [(counts[prefix], prefix)] if prefix in counts else []
35 for ch, kid in node["kids"].items():
36 self._fill(kid, prefix + ch, counts)
37 cands.extend(kid["top"])
38 node["top"] = heapq.nlargest(K, cands)
39
40 def suggest(self, prefix):
41 node = self.root
42 for ch in prefix:
43 node = node["kids"].get(ch)
44 if node is None:
45 return []
46 return [q for _, q in node["top"]]
47
48
49# ---- 1b. a flat prefix table (what you'd store in Redis or a KV store) ---------------
50
51def prefix_table(counts):
52 best = defaultdict(list)
53 for q, c in counts.items():
54 for i in range(1, len(q) + 1):
55 heap = best[q[:i]]
56 if len(heap) < K:
57 heapq.heappush(heap, (c, q))
58 elif (c, q) > heap[0]:
59 heapq.heapreplace(heap, (c, q))
60 return {p: [q for _, q in sorted(h, reverse=True)] for p, h in best.items()}
61
62
63# ---- 3. trending ------------------------------------------------------------------
64
65def decayed(count_by_hour, now_hour, half_life_h):
66 """Sum of counts, each weighted by 0.5 ** (age / half_life)."""
67 return sum(c * 0.5 ** ((now_hour - h) / half_life_h) for h, c in count_by_hour.items())
68
69
70def blended(long_term, recent, w=10):
71 return long_term + w * recent
72
73
74# ---- 4. caching short prefixes -----------------------------------------------------
75
76def share_of_requests_by_prefix_len(queries, max_len):
77 """Each search sends one request per typed character (after debounce, ~ up to the full query)."""
78 total = by_len = 0
79 for q in queries:
80 for i in range(1, len(q) + 1):
81 total += 1
82 by_len += i <= max_len
83 return by_len / total
84
85
86# ---- checks ------------------------------------------------------------------------
87
88def make_log(n_distinct=20_000, n_searches=200_000, seed=34):
89 rnd = random.Random(seed)
90 words = ["harry", "potter", "dune", "data", "database", "design", "dragon", "deep", "learning",
91 "india", "history", "kerala", "recipes", "cricket", "python", "programming", "system",
92 "interview", "atomic", "habits", "the", "art", "of", "war", "love", "story"]
93 distinct = list({" ".join(rnd.sample(words, rnd.randint(1, 3))) for _ in range(n_distinct)})
94 weights = [1 / (i + 1) ** 1.1 for i in range(len(distinct))]
95 return Counter(rnd.choices(distinct, weights=weights, k=n_searches)), distinct
96
97
98def run_checks():
99 counts, distinct = make_log()
100
101 # 1. trie and flat table agree
102 trie = Trie()
103 trie.build(counts)
104 table = prefix_table(counts)
105 for p in ["d", "da", "dat", "har", "the art", "kerala r", "zz"]:
106 assert trie.suggest(p) == table.get(p, []), p
107 top_d = table["d"]
108 assert all(q.startswith("d") for q in top_d) and len(top_d) == 5
109
110 # 2. memory for the flat table: prefixes x K suggestions x ~25 bytes each
111 n_prefixes = len(table)
112 approx_bytes = n_prefixes * K * 25
113 assert n_prefixes < sum(len(q) for q in counts) # prefixes are shared
114 # scaled-up Step 2 numbers: 50M distinct queries, ~6 new prefixes each after sharing
115 assert 50_000_000 * 6 * 5 * 25 / 1e9 == 37.5 # ~37.5 GB, before compression
116
117 # 3. trending: a new book launches; yesterday it didn't exist
118 steady = {h: 1_000 for h in range(0, 48)} # "dune": 1,000 an hour, steady
119 launch = {46: 3_000, 47: 6_000} # "dune part three": just launched
120 s_long = decayed(steady, 47, half_life_h=168) # a week-long memory
121 s_recent = decayed(steady, 47, half_life_h=2)
122 l_long = decayed(launch, 47, half_life_h=168)
123 l_recent = decayed(launch, 47, half_life_h=2)
124 assert s_long > l_long # long-term: old favourite wins
125 assert blended(l_long, l_recent, w=10) > blended(s_long, s_recent, w=10) # blended: the launch surfaces (90k vs 78k)
126
127 # 4. caching: short prefixes carry a big share of requests
128 sample = [q for q, c in counts.items() for _ in range(min(c, 50))]
129 share2 = share_of_requests_by_prefix_len(sample, 2)
130 share4 = share_of_requests_by_prefix_len(sample, 4)
131 assert 0.1 < share2 < share4 < 0.5
132 one_two = {p for p in table if len(p) <= 2}
133 assert len(one_two) < 1_000 # a few hundred keys: cache them all
134
135 # Step 2 arithmetic
136 rps = 10_000_000 * 3 * 8 / 86_400
137 assert round(rps) == 2_778 and round(rps * 5) == 13_889
138
139 return {"prefixes": n_prefixes, "approx_table_kb": approx_bytes // 1000, "top_d": top_d,
140 "share_len<=2": round(share2, 3), "share_len<=4": round(share4, 3), "one_two_keys": len(one_two)}
141
142
143if __name__ == "__main__":
144 print("All checks pass.", run_checks())