1"""Chapter 20: Probabilistic data structures. Tested with Python 3, standard library only.
2
3Running example: the bookstore.
4- Bloom filter: "is this username already taken?" without asking the database every time.
5- HyperLogLog: unique visitors today.
6- Count-min sketch + top-K: how often each search term was searched; what's trending.
7"""
8import hashlib
9import heapq
10import math
11import random
12from collections import Counter
13
14
15def hashes(item, k, m):
16 """k bucket positions from one SHA-256 (double hashing: h1 + i*h2)."""
17 d = hashlib.sha256(str(item).encode()).digest()
18 h1 = int.from_bytes(d[:8], "big")
19 h2 = int.from_bytes(d[8:16], "big") | 1
20 return [(h1 + i * h2) % m for i in range(k)]
21
22
23# ---------- 20.7 sizing ----------
24
25def bloom_size(n, p):
26 """Bits m and hash count k for n items at false-positive rate p."""
27 m = math.ceil(-n * math.log(p) / (math.log(2) ** 2))
28 k = max(1, round(m / n * math.log(2)))
29 return m, k
30
31
32def bloom_fp_rate(m, n, k):
33 return (1 - math.exp(-k * n / m)) ** k
34
35
36# ---------- 20.2 Bloom filter ----------
37
38class BloomFilter:
39 def __init__(self, n, p):
40 self.m, self.k = bloom_size(n, p)
41 self.bits = bytearray((self.m + 7) // 8)
42
43 def add(self, item):
44 for h in hashes(item, self.k, self.m):
45 self.bits[h // 8] |= 1 << (h % 8)
46
47 def might_contain(self, item):
48 return all(self.bits[h // 8] >> (h % 8) & 1 for h in hashes(item, self.k, self.m))
49
50
51# ---------- 20.3 counting Bloom filter ----------
52
53class CountingBloomFilter:
54 def __init__(self, n, p):
55 self.m, self.k = bloom_size(n, p)
56 self.counts = [0] * self.m # real ones use ~4-bit counters
57
58 def add(self, item):
59 for h in hashes(item, self.k, self.m):
60 self.counts[h] += 1
61
62 def remove(self, item):
63 if not self.might_contain(item):
64 return False
65 for h in hashes(item, self.k, self.m):
66 self.counts[h] -= 1
67 return True
68
69 def might_contain(self, item):
70 return all(self.counts[h] > 0 for h in hashes(item, self.k, self.m))
71
72
73# ---------- 20.4 HyperLogLog ----------
74
75class HyperLogLog:
76 def __init__(self, p=14):
77 self.p = p
78 self.m = 1 << p # 16,384 registers, like Redis
79 self.reg = bytearray(self.m)
80
81 def add(self, item):
82 x = int.from_bytes(hashlib.sha256(str(item).encode()).digest()[:8], "big")
83 idx = x >> (64 - self.p) # first p bits pick a register
84 rest = x & ((1 << (64 - self.p)) - 1)
85 rank = (64 - self.p) - rest.bit_length() + 1 # position of the first 1-bit
86 if rank > self.reg[idx]:
87 self.reg[idx] = rank
88
89 def count(self):
90 alpha = 0.7213 / (1 + 1.079 / self.m)
91 est = alpha * self.m * self.m / sum(2.0 ** -r for r in self.reg)
92 zeros = self.reg.count(0)
93 if est <= 2.5 * self.m and zeros: # small range: linear counting
94 est = self.m * math.log(self.m / zeros)
95 return round(est)
96
97 def merge(self, other):
98 out = HyperLogLog(self.p)
99 out.reg = bytearray(max(a, b) for a, b in zip(self.reg, other.reg))
100 return out
101
102
103# ---------- 20.5 count-min sketch ----------
104
105class CountMinSketch:
106 def __init__(self, eps, delta):
107 self.w = math.ceil(math.e / eps)
108 self.d = math.ceil(math.log(1 / delta))
109 self.table = [[0] * self.w for _ in range(self.d)]
110
111 def _cols(self, item):
112 return hashes(item, self.d, self.w)
113
114 def add(self, item, n=1):
115 for row, col in enumerate(self._cols(item)):
116 self.table[row][col] += n
117
118 def estimate(self, item):
119 return min(self.table[row][col] for row, col in enumerate(self._cols(item)))
120
121
122# ---------- 20.6 top-K ----------
123
124class TopK:
125 """Count-min sketch for counts + a small heap of the current leaders."""
126
127 def __init__(self, k, eps=0.001, delta=0.01):
128 self.k, self.cms, self.heap = k, CountMinSketch(eps, delta), {}
129
130 def add(self, item):
131 self.cms.add(item)
132 est = self.cms.estimate(item)
133 if item in self.heap or len(self.heap) < self.k:
134 self.heap[item] = est
135 else:
136 low = min(self.heap, key=self.heap.get)
137 if est > self.heap[low]:
138 del self.heap[low]
139 self.heap[item] = est
140
141 def top(self):
142 return [w for w, _ in sorted(self.heap.items(), key=lambda x: -x[1])]
143
144
145# ---------- Level 4 ----------
146
147class BuggyUsernameCheck:
148 """Treats 'might contain' as 'definitely taken'."""
149
150 def __init__(self, bloom):
151 self.bloom = bloom
152
153 def is_taken(self, name, db):
154 return self.bloom.might_contain(name) # bug: never checks the database
155
156
157def fixed_is_taken(bloom, name, db):
158 if not bloom.might_contain(name):
159 return False # "definitely not": skip the database
160 return name in db # "maybe": confirm with the truth
161
162
163# ---------- checks ----------
164
165def run_checks():
166 # sizing: 10 million usernames at 1% false positives
167 m, k = bloom_size(10_000_000, 0.01)
168 assert k == 7 and 95_000_000 < m < 96_000_000
169 mb = m / 8 / 1_000_000
170 assert 11.9 < mb < 12.0 # about 12 MB
171 assert abs(m / 10_000_000 - 9.59) < 0.01 # about 9.6 bits per item
172
173 # Bloom filter: no false negatives, ~1% false positives
174 bf = BloomFilter(20_000, 0.01)
175 taken = {f"user{i}" for i in range(20_000)}
176 for name in taken:
177 bf.add(name)
178 assert all(bf.might_contain(n) for n in taken) # never says "no" to a real one
179 fp = sum(bf.might_contain(f"new{i}") for i in range(20_000)) / 20_000
180 assert 0.005 < fp < 0.015
181 fp_overfull = bloom_fp_rate(bf.m, 60_000, bf.k) # three times the planned items
182 assert fp_overfull > 0.2
183
184 # counting Bloom filter supports delete
185 cbf = CountingBloomFilter(1_000, 0.01)
186 cbf.add("asha")
187 cbf.add("ravi")
188 assert cbf.might_contain("asha")
189 cbf.remove("asha")
190 assert not cbf.might_contain("asha") and cbf.might_contain("ravi")
191
192 # HyperLogLog: 200,000 unique visitors, many repeat visits, 16 KB of registers
193 hll_mon, hll_tue = HyperLogLog(), HyperLogLog()
194 rnd = random.Random(20)
195 monday = [rnd.randrange(150_000) for _ in range(300_000)] # 300k visits, repeats included
196 for v in monday:
197 hll_mon.add(f"visitor{v}")
198 for i in range(100_000, 200_000):
199 hll_tue.add(f"visitor{i}")
200 true_mon = len(set(monday))
201 est_mon = hll_mon.count()
202 assert abs(est_mon - true_mon) / true_mon < 0.02
203 both = hll_mon.merge(hll_tue).count() # union without the raw lists
204 true_both = len(set(monday) | set(range(100_000, 200_000)))
205 assert abs(both - true_both) / true_both < 0.02
206 assert round(1.04 / math.sqrt(16_384), 4) == 0.0081
207
208 # count-min sketch on a skewed stream of search terms
209 terms = [f"term{int(rnd.paretovariate(1.1))}" for _ in range(100_000)]
210 truth = Counter(terms)
211 cms = CountMinSketch(eps=0.001, delta=0.01)
212 for t in terms:
213 cms.add(t)
214 assert cms.w == 2_719 and cms.d == 5
215 assert all(cms.estimate(t) >= c for t, c in truth.items()) # never under
216 over = [cms.estimate(t) - c for t, c in truth.items()]
217 assert sum(o <= 0.001 * len(terms) for o in over) / len(over) > 0.99
218
219 # top-K matches the true leaders
220 tk = TopK(5)
221 for t in terms:
222 tk.add(t)
223 assert tk.top() == [t for t, _ in truth.most_common(5)]
224
225 # Level 4: the buggy check rejects free names
226 db = taken
227 buggy = BuggyUsernameCheck(bf)
228 wrongly_taken = [f"new{i}" for i in range(20_000) if buggy.is_taken(f"new{i}", db)]
229 assert len(wrongly_taken) > 100
230 assert not any(fixed_is_taken(bf, n, db) for n in wrongly_taken)
231 assert fixed_is_taken(bf, "user42", db)
232
233 # exact-set comparison (Level 2)
234 exact_bytes = 10_000_000 * 20 # ~20 bytes per username, ignoring overhead
235 assert exact_bytes / 1_000_000 == 200
236
237 return {"bloom_mb": round(mb, 2), "fp": fp, "fp_overfull": round(fp_overfull, 2),
238 "hll": (est_mon, true_mon), "hll_union": (both, true_both),
239 "cms": (cms.w, cms.d), "top": tk.top(), "wrongly_taken": len(wrongly_taken)}
240
241
242if __name__ == "__main__":
243 print("All checks pass.", run_checks())