129 lines
4.4 KiB
Python
129 lines
4.4 KiB
Python
"""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)
|