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