diff --git a/wizard-work/beam/sparse_table.py b/wizard-work/beam/sparse_table.py new file mode 100644 index 0000000..fb90fc5 --- /dev/null +++ b/wizard-work/beam/sparse_table.py @@ -0,0 +1,45 @@ +"""Sparse table for O(1) range queries over an idempotent associative op. + +Idempotent means f(x, x) == x (min, max, gcd, bitwise AND/OR). For such ops +the overlap trick gives O(1) queries after O(n log n) preprocessing. + +Usage: + st = SparseTable([3,1,4,1,5,9,2,6]) + st.query(1, 4) # min of positions 1..4 (inclusive, 0-indexed) -> 1 +""" +from __future__ import annotations +import math +from typing import Callable, Sequence + + +class SparseTable: + def __init__(self, data: Sequence, func: Callable = min): + self.n = len(data) + self.func = func + self._log = [0] * (self.n + 1) + for i in range(2, self.n + 1): + self._log[i] = self._log[i // 2] + 1 + self.k = self._log[self.n] if self.n > 0 else 0 + # st[j][i] = func over data[i .. i + 2^j - 1] + self.st = [] + if self.n > 0: + self.st.append(list(data)) + j = 1 + while (1 << j) <= self.n: + prev = self.st[j - 1] + span = 1 << (j - 1) + cur = [func(prev[i], prev[i + span]) for i in range(self.n - (1 << j) + 1)] + self.st.append(cur) + j += 1 + + def query(self, l: int, r: int): + """Inclusive range [l, r], 0-indexed. Raises on empty/invalid.""" + if self.n == 0: + raise IndexError("empty table") + if l < 0 or r >= self.n or l > r: + raise IndexError(f"invalid range [{l}, {r}] for n={self.n}") + j = self._log[r - l + 1] + return self.func(self.st[j][l], self.st[j][r - (1 << j) + 1]) + + def __len__(self): + return self.n diff --git a/wizard-work/beam/test_sparse_table.py b/wizard-work/beam/test_sparse_table.py new file mode 100644 index 0000000..3566031 --- /dev/null +++ b/wizard-work/beam/test_sparse_table.py @@ -0,0 +1,67 @@ +import random +import unittest +from sparse_table import SparseTable + + +class TestSparseTable(unittest.TestCase): + def test_min_basic(self): + a = [3, 1, 4, 1, 5, 9, 2, 6] + st = SparseTable(a) + self.assertEqual(st.query(0, 7), 1) + self.assertEqual(st.query(1, 4), 1) + self.assertEqual(st.query(2, 5), 1) + self.assertEqual(st.query(6, 7), 2) + + def test_single(self): + st = SparseTable([42]) + self.assertEqual(st.query(0, 0), 42) + + def test_empty_raises(self): + st = SparseTable([]) + with self.assertRaises(IndexError): + st.query(0, 0) + + def test_invalid_range(self): + st = SparseTable([1, 2, 3]) + for l, r in [(-1, 1), (0, 3), (2, 1)]: + with self.assertRaises(IndexError): + st.query(l, r) + + def test_max_func(self): + a = [3, 1, 4, 1, 5, 9, 2, 6] + st = SparseTable(a, func=max) + self.assertEqual(st.query(0, 7), 9) + self.assertEqual(st.query(1, 4), 5) + + def test_gcd_func(self): + import math + a = [12, 18, 24, 30] + st = SparseTable(a, func=math.gcd) + self.assertEqual(st.query(0, 3), 6) + self.assertEqual(st.query(1, 2), 6) + + def test_fuzz_min(self): + rng = random.Random(7) + for _ in range(200): + n = rng.randint(1, 60) + a = [rng.randint(-100, 100) for _ in range(n)] + st = SparseTable(a) + for _ in range(50): + l = rng.randrange(n) + r = rng.randrange(l, n) + self.assertEqual(st.query(l, r), min(a[l:r + 1])) + + def test_fuzz_max(self): + rng = random.Random(9) + for _ in range(100): + n = rng.randint(1, 40) + a = [rng.randint(-50, 50) for _ in range(n)] + st = SparseTable(a, func=max) + for _ in range(30): + l = rng.randrange(n) + r = rng.randrange(l, n) + self.assertEqual(st.query(l, r), max(a[l:r + 1])) + + +if __name__ == "__main__": + unittest.main()