Interval tree in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/interval-tree/. The authors are listed in prov/8H8C9B.jsonl.

import random
from typing import Optional


class _Node:
    __slots__ = ("lo", "hi", "uid", "prio", "left", "right", "max")

    def __init__(self, lo: int, hi: int, uid: int, prio: float) -> None:
        self.lo = lo
        self.hi = hi
        self.uid = uid
        self.prio = prio
        self.left: Optional["_Node"] = None
        self.right: Optional["_Node"] = None
        self.max = hi


def _mx(x: Optional[_Node]) -> float:
    return x.max if x is not None else float("-inf")


def _pull(x: _Node) -> _Node:
    x.max = max(x.hi, _mx(x.left), _mx(x.right))
    return x


def _key(x: _Node) -> tuple[int, int, int]:
    return (x.lo, x.hi, x.uid)


def _split(t: Optional[_Node], key: tuple[int, int, int]) -> tuple[Optional[_Node], Optional[_Node]]:
    """Split t into the nodes with key < key and those with key >= key."""
    if t is None:
        return None, None
    if _key(t) < key:
        a, b = _split(t.right, key)
        t.right = a
        return _pull(t), b
    a, b = _split(t.left, key)
    t.left = b
    return a, _pull(t)


def _merge(a: Optional[_Node], b: Optional[_Node]) -> Optional[_Node]:
    """Merge two treaps where all keys of a are below all keys of b."""
    if a is None:
        return b
    if b is None:
        return a
    if a.prio > b.prio:
        a.right = _merge(a.right, b)
        return _pull(a)
    b.left = _merge(a, b.left)
    return _pull(b)


class IntervalTree:
    """A dynamic multiset of closed intervals [lo, hi] in a treap keyed by lo,
    with the largest right end of each subtree in the field max.

    insert, delete and search take O(log n) expected time; report takes
    O((k + 1) log n) expected time for k results. Space O(n).
    """

    def __init__(self, seed: int = 0) -> None:
        self.root: Optional[_Node] = None
        self._rng = random.Random(seed)
        self._next = 0
        self.size = 0

    def insert(self, lo: int, hi: int) -> None:
        """Add the interval [lo, hi]; requires lo <= hi."""
        if lo > hi:
            raise ValueError("lo must be at most hi")
        node = _Node(lo, hi, self._next, self._rng.random())
        self._next += 1
        a, b = _split(self.root, _key(node))
        self.root = _merge(_merge(a, node), b)
        self.size += 1

    def delete(self, lo: int, hi: int) -> bool:
        """Remove one copy of [lo, hi]; return False if there is none."""
        a, rest = _split(self.root, (lo, hi, -1))
        mid, b = _split(rest, (lo, hi + 1, -1))
        found = mid is not None
        if mid is not None:
            # every node of mid holds [lo, hi]; drop its root
            mid = _merge(mid.left, mid.right)
            self.size -= 1
        self.root = _merge(_merge(a, mid), b)
        return found

    def search(self, a: int, b: int) -> Optional[tuple[int, int]]:
        """One stored interval that overlaps [a, b], or None. O(height)."""
        x = self.root  # @line L1
        while x is not None and not (x.lo <= b and a <= x.hi):  # @line L2
            if x.left is not None and x.left.max >= a:  # @line L3
                x = x.left
            else:  # @line L4
                x = x.right
        return (x.lo, x.hi) if x is not None else None  # @line L5

    def report(self, a: int, b: int) -> list[tuple[int, int]]:
        """All stored intervals that overlap [a, b], in order of (lo, hi)."""
        out: list[tuple[int, int]] = []
        stack: list[tuple[Optional[_Node], bool]] = [(self.root, False)]
        while stack:
            x, expanded = stack.pop()
            if x is None:
                continue
            if not expanded:
                if x.max < a:  # @line L6
                    continue
                if x.lo <= b:  # @line L9
                    stack.append((x.right, False))
                stack.append((x, True))
                stack.append((x.left, False))  # @line L7
            elif x.lo <= b and a <= x.hi:  # @line L8
                out.append((x.lo, x.hi))
        return out

    def check_invariant(self) -> None:
        """Check the search order, the heap order of priorities, the max
        fields and the size."""
        count = 0
        stack: list[tuple[Optional[_Node], Optional[tuple[int, int, int]], Optional[tuple[int, int, int]]]] = [
            (self.root, None, None)
        ]
        while stack:
            x, low, high = stack.pop()
            if x is None:
                continue
            count += 1
            k = _key(x)
            assert low is None or low < k, "search order"
            assert high is None or k < high, "search order"
            assert x.lo <= x.hi, "empty interval"
            for c in (x.left, x.right):
                if c is not None:
                    assert c.prio <= x.prio, "heap order"
            assert x.max == max(x.hi, _mx(x.left), _mx(x.right)), "max field"
            stack.append((x.left, low, k))
            stack.append((x.right, k, high))
        assert count == self.size, "size"