beam: add SparseTable (O(1) idempotent range query) + randomized tests

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

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()