Flajolet-Martin sketch in Python
Estimates distinct count from the maximum number of trailing zeros among hashed values.
- File
flajolet_martin_sketch.py, 64 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/flajolet-martin-sketch/. The authors are listed in prov/86TWSR.jsonl.
"""Flajolet-Martin probabilistic counting with stochastic averaging (PCSA, 1985)."""
import hashlib
Key = str | bytes | int
PHI = 0.77351 # the correction constant of Flajolet and Martin
def rho(q: int) -> int:
"""The number of trailing zero bits of q > 0."""
return (q & -q).bit_length() - 1
class FMSketch:
"""m bitmaps of L bits. add is O(1); estimate is O(m L).
The state depends only on the set of distinct keys added, so duplicates
and the order of the stream do not matter.
"""
def __init__(self, m: int = 64, L: int = 32, seed: int = 0) -> None:
if m < 1 or L < 1:
raise ValueError("m and L must be positive")
self.m, self.L, self.seed = m, L, seed
self.bitmaps = [0] * m # @line L1: bit r of bitmaps[j] is B_j[r]
def _hash(self, key: Key) -> int:
d = hashlib.blake2b(repr(key).encode(), digest_size=8, key=self.seed.to_bytes(8, "little")) # @line L2
return int.from_bytes(d.digest(), "little")
def add(self, key: Key) -> None:
"""Record one occurrence of key."""
y = self._hash(key)
j, q = y % self.m, y // self.m # @line L3
if q:
r = rho(q)
if r < self.L:
self.bitmaps[j] |= 1 << r # @line L4
def rank(self, j: int) -> int:
"""R_j: the index of the lowest 0 bit of bitmap j."""
b, r = self.bitmaps[j], 0
while r < self.L and b >> r & 1:
r += 1
return r
def estimate(self) -> float:
"""The PCSA estimate (m / PHI) * 2^(mean of the R_j)."""
s = 0 # @line L5
for j in range(self.m):
s += self.rank(j) # @line L6
return self.m / PHI * 2 ** (s / self.m) # @line L7
def merge(self, other: "FMSketch") -> "FMSketch":
"""The sketch of the union of the two key sets."""
if (self.m, self.L, self.seed) != (other.m, other.L, other.seed):
raise ValueError("sketches must have the same m, L and seed")
out = FMSketch(self.m, self.L, self.seed)
out.bitmaps = [a | b for a, b in zip(self.bitmaps, other.bitmaps)]
return out
test_flajolet_martin_sketch.py 82 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/flajolet-martin-sketch/. The authors are listed in prov/86TWSR.jsonl.
import random
import pytest
from flajolet_martin_sketch import PHI, FMSketch, Key, rho
HASHES: dict[Key, int] = {"a": 12, "b": 7, "c": 4, "d": 16, "e": 9, "f": 2, "g": 11}
class FixedFM(FMSketch):
def _hash(self, key: Key) -> int:
return HASHES[key]
def test_example() -> None:
s = FixedFM(m=2, L=8)
for k in "abcadefgb":
s.add(k)
assert s.bitmaps == [0b1011, 0b101]
assert s.rank(0) == 2 and s.rank(1) == 1
assert s.estimate() == pytest.approx(2 / PHI * 2 ** 1.5)
assert round(s.estimate(), 2) == 7.31
def test_edges() -> None:
assert rho(1) == 0 and rho(8) == 3 and rho(12) == 2
s = FMSketch(m=1, L=4)
assert s.rank(0) == 0
s.bitmaps[0] = 0b1111
assert s.rank(0) == 4
with pytest.raises(ValueError):
FMSketch(m=0)
with pytest.raises(ValueError):
FMSketch(m=2).merge(FMSketch(m=4))
def brute(keys: set[int], m: int, L: int, seed: int) -> list[int]:
s = FMSketch(m, L, seed)
bm = [0] * m
for x in keys:
y = s._hash(x)
j, q = y % m, y // m
for r in range(L):
if q and q % (1 << (r + 1)) == (1 << r):
bm[j] |= 1 << r
return bm
def test_random_set_function() -> None:
rng = random.Random(1985)
for _ in range(500):
m, L, seed = rng.randint(1, 8), rng.randint(1, 12), rng.randrange(50)
keys = {rng.randrange(1000) for _ in range(rng.randint(0, 40))}
stream = [x for x in keys for _ in range(rng.randint(1, 3))]
rng.shuffle(stream)
a = FMSketch(m, L, seed)
for x in stream:
a.add(x)
assert a.bitmaps == brute(keys, m, L, seed)
part = set(rng.sample(sorted(keys), len(keys) // 2))
b, c = FMSketch(m, L, seed), FMSketch(m, L, seed)
for x in part:
b.add(x)
for x in keys - part:
c.add(x)
assert b.merge(c).bitmaps == a.bitmaps
assert b.merge(c).estimate() == a.estimate()
def test_accuracy() -> None:
errs = []
for seed in range(10):
s = FMSketch(m=64, L=32, seed=seed)
n = 20000
for i in range(n):
s.add(i)
errs.append(s.estimate() / n - 1)
assert all(abs(e) < 0.4 for e in errs)
assert abs(sum(errs) / len(errs)) < 0.15
What the next level needs
- the tests did not run in this build