Disjoint sparse table in Python
O(n log n) table answering any associative range query in O(1) using divide and conquer.
- File
disjoint_sparse_table.py, 63 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/disjoint-sparse-table/. The authors are listed in prov/5HH6N4.jsonl.
from typing import Callable, Generic, Sequence, TypeVar
T = TypeVar("T")
class DisjointSparseTable(Generic[T]):
"""Range products A[l] op ... op A[r] for any associative op.
Build: O(n log n) calls of op and n*ceil(lg n) stored values.
Query: O(1) time and at most one call of op. The operation need not be
commutative, idempotent or invertible.
"""
def __init__(self, a: Sequence[T], op: Callable[[T, T], T]) -> None:
"""Build the table of a with the associative operation op."""
self.a = list(a)
self.op = op
n = len(self.a)
self.n = n
self.k = max(0, (n - 1).bit_length())
self.table: list[list[T]] = [self.a] # level 0 is the array
for h in range(1, self.k + 1): # @line L4
row = list(self.a)
size = 1 << h
for b in range(0, n, size): # @line L5
m = b + size // 2 # @line L6
if m - 1 < n: # @line L7
row[m - 1] = self.a[m - 1]
for i in range(min(m, n) - 2, b - 1, -1): # @line L8
row[i] = op(self.a[i], row[i + 1])
if m < n: # @line L9
row[m] = self.a[m]
for i in range(m + 1, min(b + size, n)): # @line L10
row[i] = op(row[i - 1], self.a[i])
self.table.append(row)
def query(self, l: int, r: int) -> T:
"""Return A[l] op ... op A[r] for 0 <= l <= r < n."""
if not 0 <= l <= r < self.n:
raise IndexError("need 0 <= l <= r < n")
if l == r: # @line L1
return self.a[l]
h = (l ^ r).bit_length() # @line L2
return self.op(self.table[h][l], self.table[h][r]) # @line L3
def check_invariant(self) -> None:
"""Recompute each entry from its definition and compare."""
n = self.n
for h in range(1, self.k + 1):
size = 1 << h
for i in range(n):
b = (i // size) * size
m = b + size // 2
lo, hi = (i, m - 1) if i < m else (m, i)
if hi >= n:
continue
acc = self.a[lo]
for j in range(lo + 1, hi + 1):
acc = self.op(acc, self.a[j])
assert self.table[h][i] == acc, "table entry"
test_disjoint_sparse_table.py 64 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/disjoint-sparse-table/. The authors are listed in prov/5HH6N4.jsonl.
import operator
import random
import pytest
from disjoint_sparse_table import DisjointSparseTable
Mat = tuple[int, int, int, int]
MOD = 1_000_003
def matmul(x: Mat, y: Mat) -> Mat:
a, b, c, d = x
e, f, g, h = y
return ((a * e + b * g) % MOD, (a * f + b * h) % MOD, (c * e + d * g) % MOD, (c * f + d * h) % MOD)
def test_example() -> None:
a = [3, 1, 4, 1, 5, 9, 2, 6]
t = DisjointSparseTable(a, operator.add)
assert t.table[1] == [3, 1, 4, 1, 5, 9, 2, 6]
assert t.table[2] == [4, 1, 4, 5, 14, 9, 2, 8]
assert t.table[3] == [9, 6, 5, 1, 5, 14, 16, 22]
assert t.query(2, 5) == 19
assert t.query(4, 7) == 22
assert t.query(1, 2) == 5
assert t.query(3, 3) == 1
t.check_invariant()
def test_edges() -> None:
t = DisjointSparseTable(["x"], operator.add)
assert t.query(0, 0) == "x"
with pytest.raises(IndexError):
t.query(0, 1)
e = DisjointSparseTable([], operator.add)
with pytest.raises(IndexError):
e.query(0, 0)
s = DisjointSparseTable(list("abcde"), operator.add)
assert s.query(0, 4) == "abcde" and s.query(1, 3) == "bcd" and s.query(3, 4) == "de"
s.check_invariant()
def test_random_against_brute_force() -> None:
rng = random.Random(1982)
for _ in range(500):
n = rng.randint(1, 40)
mats = [tuple(rng.randrange(7) for _ in range(4)) for _ in range(n)]
t = DisjointSparseTable(mats, matmul)
words = ["".join(rng.choice("ab") for _ in range(rng.randint(0, 2))) for _ in range(n)]
u = DisjointSparseTable(words, operator.add)
for _ in range(10):
l = rng.randrange(n)
r = rng.randrange(l, n)
acc = mats[l]
for j in range(l + 1, r + 1):
acc = matmul(acc, mats[j])
assert t.query(l, r) == acc
assert u.query(l, r) == "".join(words[l : r + 1])
if n <= 20:
t.check_invariant()
What the next level needs
- the tests did not run in this build