beam: add interval_tree with point-stabbing queries + tests

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)

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