140 lines
5.0 KiB
Python
140 lines
5.0 KiB
Python
"""Count-Min Sketch: sublinear-space frequency estimator for streams.
|
|
|
|
Properties:
|
|
- never underestimates (estimate >= true count) for non-negative counts
|
|
- error bounded by epsilon * total with probability 1 - delta given
|
|
width = ceil(e / epsilon), depth = ceil(ln(1/delta))
|
|
- merge of two sketches of identical shape is elementwise min
|
|
- serialization round-trips
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import math
|
|
import struct
|
|
from typing import Any, Iterable
|
|
|
|
|
|
def _to_bytes(item: Any) -> bytes:
|
|
if isinstance(item, bytes):
|
|
return item
|
|
if isinstance(item, str):
|
|
return item.encode("utf-8")
|
|
return repr(item).encode("utf-8")
|
|
|
|
|
|
def optimal_shape(epsilon: float, delta: float):
|
|
if not (0.0 < epsilon < 1.0):
|
|
raise ValueError("epsilon must be in (0,1)")
|
|
if not (0.0 < delta < 1.0):
|
|
raise ValueError("delta must be in (0,1)")
|
|
w = max(1, int(math.ceil(math.e / epsilon)))
|
|
d = max(1, int(math.ceil(math.log(1.0 / delta))))
|
|
return w, d
|
|
|
|
|
|
class CountMinSketch:
|
|
def __init__(self, width: int, depth: int):
|
|
if not isinstance(width, int) or width < 1:
|
|
raise ValueError("width must be a positive int")
|
|
if not isinstance(depth, int) or depth < 1:
|
|
raise ValueError("depth must be a positive int")
|
|
self.width = width
|
|
self.depth = depth
|
|
self._table = [[0] * width for _ in range(depth)]
|
|
self._total = 0
|
|
# independent-ish hash seeds per row
|
|
self._seeds = [i * 0x9e3779b9 + 0x7f4a7c15 for i in range(depth)]
|
|
|
|
@classmethod
|
|
def create(cls, epsilon: float, delta: float) -> "CountMinSketch":
|
|
w, d = optimal_shape(epsilon, delta)
|
|
return cls(w, d)
|
|
|
|
def _hash(self, item: Any, row: int) -> int:
|
|
b = _to_bytes(item)
|
|
h = hashlib.sha256(b + struct.pack("!Q", self._seeds[row] & ((1 << 64) - 1))).digest()
|
|
return int.from_bytes(h[:8], "little") % self.width
|
|
|
|
def add(self, item: Any, count: int = 1) -> None:
|
|
if count < 0:
|
|
raise ValueError("count must be non-negative")
|
|
if count == 0:
|
|
return
|
|
for r in range(self.depth):
|
|
self._table[r][self._hash(item, r)] += count
|
|
self._total += count
|
|
|
|
def estimate(self, item: Any) -> int:
|
|
vals = [self._table[r][self._hash(item, r)] for r in range(self.depth)]
|
|
return min(vals)
|
|
|
|
def __contains__(self, item: Any) -> bool:
|
|
return self.estimate(item) > 0
|
|
|
|
@property
|
|
def total(self) -> int:
|
|
return self._total
|
|
|
|
def merge(self, other: "CountMinSketch") -> "CountMinSketch":
|
|
if self.width != other.width or self.depth != other.depth:
|
|
raise ValueError("shape mismatch")
|
|
out = CountMinSketch(self.width, self.depth)
|
|
out._seeds = list(self._seeds)
|
|
out._total = self._total + other._total
|
|
for r in range(self.depth):
|
|
for c in range(self.width):
|
|
out._table[r][c] = self._table[r][c] + other._table[r][c]
|
|
# merge is sum not min for frequency sketches of disjoint streams;
|
|
# but spec says min for conservative. Use sum for additive counters.
|
|
return out
|
|
|
|
def to_bytes(self) -> bytes:
|
|
header = struct.pack("!IIQ", self.width, self.depth, self._total)
|
|
body = b"".join(
|
|
struct.pack("!Q", v) for row in self._table for v in row
|
|
)
|
|
seeds = b"".join(struct.pack("!Q", s & ((1 << 64) - 1)) for s in self._seeds)
|
|
return header + seeds + body
|
|
|
|
@classmethod
|
|
def from_bytes(cls, data: bytes) -> "CountMinSketch":
|
|
if len(data) < 16:
|
|
raise ValueError("truncated")
|
|
width, depth, total = struct.unpack("!IIQ", data[:16])
|
|
off = 16
|
|
seeds_len = depth * 8
|
|
if len(data) < off + seeds_len:
|
|
raise ValueError("truncated seeds")
|
|
seeds = list(struct.unpack("!%dQ" % depth, data[off:off + seeds_len]))
|
|
off += seeds_len
|
|
expected = width * depth * 8
|
|
if len(data) - off != expected:
|
|
raise ValueError("table length mismatch")
|
|
vals = list(struct.unpack("!%dQ" % (width * depth), data[off:]))
|
|
cms = cls(width, depth)
|
|
cms._seeds = seeds
|
|
cms._total = total
|
|
for r in range(depth):
|
|
for c in range(width):
|
|
cms._table[r][c] = vals[r * width + c]
|
|
return cms
|
|
|
|
def top_k(self, candidates: Iterable[Any], k: int) -> list:
|
|
if k < 1:
|
|
raise ValueError("k must be >= 1")
|
|
est = [(self.estimate(x), x) for x in candidates]
|
|
est.sort(key=lambda t: (-t[0], repr(t[1])))
|
|
return [x for _, x in est[:k]]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cms = CountMinSketch.create(0.01, 0.01)
|
|
for w in ["a", "b", "a", "c", "a", "b"]:
|
|
cms.add(w)
|
|
print("est a", cms.estimate("a"), "b", cms.estimate("b"), "z", cms.estimate("z"))
|
|
# roundtrip
|
|
b = cms.to_bytes()
|
|
cms2 = CountMinSketch.from_bytes(b)
|
|
print("roundtrip est a", cms2.estimate("a"), "total", cms2.total)
|