From 0dce52d3bb87dabeda01b58d684175ebc954d7ab Mon Sep 17 00:00:00 2001 From: beam Date: Sat, 31 Oct 2026 10:12:53 +0000 Subject: [PATCH] beam: add interval_tree with point-stabbing queries + tests --- wizard-work/beam/interval_tree.py | 128 +++++++++++++++++++++++++ wizard-work/beam/test_interval_tree.py | 83 ++++++++++++++++ 2 files changed, 211 insertions(+) create mode 100644 wizard-work/beam/interval_tree.py create mode 100644 wizard-work/beam/test_interval_tree.py diff --git a/wizard-work/beam/interval_tree.py b/wizard-work/beam/interval_tree.py new file mode 100644 index 0000000..f38f73d --- /dev/null +++ b/wizard-work/beam/interval_tree.py @@ -0,0 +1,128 @@ +"""Interval tree (center-point partitioning) with point-stabbing queries. + +An interval tree stores closed intervals [low, high] and answers: + which intervals contain point x? +in O(log n + k) where k is the number of reported intervals. + +Design +------ +Each node holds: + - center: a point (median of midpoints of the node's intervals) + - left: intervals entirely to the left of center (high < center) + - right: intervals entirely to the right of center (low > center) + - cross: intervals crossing center (low <= center <= high), stored twice: + cross_by_low sorted by low (ascending) + cross_by_high sorted by high (descending) + - left_child, right_child + +Query at point x: + - If x == center: every interval in cross contains x. + - If x < center: intervals in cross with low <= x contain x + (since high >= center > x automatically); recurse left. + - If x > center: intervals in cross with high >= x contain x + (since low <= center < x automatically); recurse right. + +Insert/delete rebuild the affected subtree from its interval list, keeping +the structure simple and obviously correct. Complexity is O(n log n) per +mutating op in the worst case, acceptable for the intended use (correctness +over micro-performance). +""" + +from __future__ import annotations +from dataclasses import dataclass, field +from typing import Iterable, List, Tuple + +Interval = Tuple[int, int] + + +@dataclass +class _Node: + center: int + cross_by_low: List[Interval] = field(default_factory=list) + cross_by_high: List[Interval] = field(default_factory=list) + left_child: "_Node | None" = None + right_child: "_Node | None" = None + left_intervals: List[Interval] = field(default_factory=list) + right_intervals: List[Interval] = field(default_factory=list) + + +def _build(intervals: List[Interval]) -> "_Node | None": + if not intervals: + return None + # center = median of midpoints + mids = sorted((lo + hi) // 2 for lo, hi in intervals) + center = mids[len(mids) // 2] + cross, left, right = [], [], [] + for lo, hi in intervals: + if hi < center: + left.append((lo, hi)) + elif lo > center: + right.append((lo, hi)) + else: + cross.append((lo, hi)) + node = _Node(center=center) + node.cross_by_low = sorted(cross, key=lambda iv: iv[0]) + node.cross_by_high = sorted(cross, key=lambda iv: iv[1], reverse=True) + node.left_intervals = left + node.right_intervals = right + node.left_child = _build(left) + node.right_child = _build(right) + return node + + +class IntervalTree: + def __init__(self, intervals: Iterable[Interval] = ()): + self._intervals: List[Interval] = list(intervals) + # validate + for lo, hi in self._intervals: + if lo > hi: + raise ValueError(f"invalid interval: low {lo} > high {hi}") + self._root = _build(self._intervals) + + def _rebuild(self) -> None: + self._root = _build(self._intervals) + + def insert(self, low: int, high: int) -> None: + if low > high: + raise ValueError(f"invalid interval: low {low} > high {high}") + self._intervals.append((low, high)) + self._rebuild() + + def delete(self, low: int, high: int) -> bool: + try: + self._intervals.remove((low, high)) + except ValueError: + return False + self._rebuild() + return True + + def query_point(self, x: int) -> List[Interval]: + res: List[Interval] = [] + node = self._root + while node is not None: + if x == node.center: + res.extend(node.cross_by_low) + # no need to recurse: all cross contain center==x; + # left/right intervals are strictly separated + break + if x < node.center: + for lo, hi in node.cross_by_low: + if lo <= x: # hi >= center > x guaranteed + res.append((lo, hi)) + else: + break + node = node.left_child + else: # x > center + for lo, hi in node.cross_by_high: + if hi >= x: # lo <= center < x guaranteed + res.append((lo, hi)) + else: + break + node = node.right_child + return res + + def all_intervals(self) -> List[Interval]: + return list(self._intervals) + + def __len__(self) -> int: + return len(self._intervals) diff --git a/wizard-work/beam/test_interval_tree.py b/wizard-work/beam/test_interval_tree.py new file mode 100644 index 0000000..53c07c0 --- /dev/null +++ b/wizard-work/beam/test_interval_tree.py @@ -0,0 +1,83 @@ +import random +from interval_tree import IntervalTree + + +def test_empty(): + t = IntervalTree() + assert t.query_point(5) == [] + assert len(t) == 0 + + +def test_basic(): + t = IntervalTree([(1, 5), (2, 3), (4, 8), (10, 12)]) + assert sorted(t.query_point(3)) == [(1, 5), (2, 3)] + assert sorted(t.query_point(4)) == [(1, 5), (4, 8)] + assert sorted(t.query_point(11)) == [(10, 12)] + assert t.query_point(0) == [] + assert t.query_point(100) == [] + + +def test_point_intervals(): + t = IntervalTree([(5, 5), (5, 5), (1, 9)]) + assert sorted(t.query_point(5)) == [(1, 9), (5, 5), (5, 5)] + + +def test_insert_delete(): + t = IntervalTree() + t.insert(1, 4) + t.insert(3, 6) + assert sorted(t.query_point(3)) == [(1, 4), (3, 6)] + assert t.delete(1, 4) is True + assert t.delete(1, 4) is False + assert t.query_point(3) == [(3, 6)] + assert len(t) == 1 + + +def test_invalid(): + try: + IntervalTree([(2, 1)]) + raise AssertionError("expected ValueError") + except ValueError: + pass + t = IntervalTree() + try: + t.insert(5, 1) + raise AssertionError("expected ValueError") + except ValueError: + pass + + +def test_random_brute(): + rng = random.Random(42) + for trial in range(200): + n = rng.randint(0, 40) + ivs = [] + for _ in range(n): + a = rng.randint(-20, 20) + b = rng.randint(-20, 20) + lo, hi = (a, b) if a <= b else (b, a) + ivs.append((lo, hi)) + t = IntervalTree(ivs) + for x in range(-25, 26): + got = sorted(t.query_point(x)) + exp = sorted(iv for iv in ivs if iv[0] <= x <= iv[1]) + assert got == exp, (trial, x, got, exp) + # mutate + for _ in range(5): + if rng.random() < 0.5 and ivs: + iv = rng.choice(ivs) + t.delete(*iv) + ivs.remove(iv) + else: + a = rng.randint(-20, 20); b = rng.randint(-20, 20) + lo, hi = (a, b) if a <= b else (b, a) + t.insert(lo, hi); ivs.append((lo, hi)) + for x in range(-25, 26): + got = sorted(t.query_point(x)) + exp = sorted(iv for iv in ivs if iv[0] <= x <= iv[1]) + assert got == exp, (trial, x, got, exp) + + +if __name__ == "__main__": + import pytest, sys + sys.exit(pytest.main([__file__, "-q"]))