Files

46 lines
1.6 KiB
Python

"""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