PR quadtree in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/pr-quadtree/. The authors are listed in prov/D4SJDH.jsonl.

from typing import Optional

Point = tuple[int, int]


class Node:
    """A square cell with corner (x, y) and side size.

    A leaf has children None and holds at most B points. An internal node has
    four children, numbered o = bx + 2*by (0 SW, 1 SE, 2 NW, 3 NE), where bx is
    1 for the east half and by is 1 for the north half. count is the number of
    points in the subtree.
    """

    __slots__ = ("x", "y", "size", "points", "children", "count")

    def __init__(self, x: int, y: int, size: int) -> None:
        self.x, self.y, self.size = x, y, size
        self.points: list[Point] = []
        self.children: Optional[list[Node]] = None
        self.count = 0

    def quadrant(self, p: Point) -> int:
        """The child cell that holds p. O(1)."""
        h = self.size // 2
        return (1 if p[0] >= self.x + h else 0) | (2 if p[1] >= self.y + h else 0)

    def holds(self, p: Point) -> bool:
        return self.x <= p[0] < self.x + self.size and self.y <= p[1] < self.y + self.size


class PrQuadtree:
    """A PR quadtree with leaf capacity B over the square [0, 2^depth)^2.

    Invariant: each point lies in the cell of its leaf, a leaf holds at most B
    points, and each internal node has more than B points below it. So the
    shape depends only on the set of points. Points have integer coordinates
    and are stored as a set; the height is at most depth.
    """

    def __init__(self, depth: int, capacity: int = 1) -> None:
        if depth < 0 or capacity < 1:
            raise ValueError("need depth >= 0 and capacity >= 1")
        self.depth = depth
        self.capacity = capacity
        self.root = Node(0, 0, 1 << depth)

    def __len__(self) -> int:
        return self.root.count

    def _norm(self, p: Point) -> Point:
        q = (int(p[0]), int(p[1]))
        if not self.root.holds(q):
            raise ValueError("the point is outside the root square")
        return q

    def _leaf(self, p: Point) -> tuple[Node, list[Node]]:
        """The leaf whose cell holds p and the internal nodes above it. O(depth)."""
        u = self.root
        path: list[Node] = []
        while u.children is not None:
            path.append(u)
            u = u.children[u.quadrant(p)]
        return u, path

    def contains(self, p: Point) -> bool:
        """Whether p is stored. O(depth + B)."""
        q = self._norm(p)
        return q in self._leaf(q)[0].points

    def insert(self, p: Point) -> bool:
        """Add p; return False if it is already stored. O(depth * B) worst."""
        q = self._norm(p)
        u, path = self._leaf(q)  # @line L1
        if q in u.points:
            return False
        for v in path:  # @line L2
            v.count += 1
        u.points.append(q)  # @line L3
        u.count += 1
        B = self.capacity
        while len(u.points) > B:  # @line L4: a cell of side 1 never overflows
            h = u.size // 2
            u.children = [Node(u.x + (o & 1) * h, u.y + (o >> 1) * h, h) for o in range(4)]  # @line L5
            for r in u.points:  # @line L6
                c = u.children[u.quadrant(r)]
                c.points.append(r)
                c.count += 1
            u.points = []
            nxt = None
            for c in u.children:  # @line L7: at most one child overflows
                if len(c.points) > B:
                    nxt = c
            if nxt is None:
                break
            u = nxt
        return True  # @line L8

    def range_search(self, lo: Point, hi: Point) -> list[Point]:
        """The stored points in [lo[0], hi[0]] x [lo[1], hi[1]]; explicit stack."""
        out: list[Point] = []
        stack = [self.root]
        while stack:
            u = stack.pop()
            if hi[0] < u.x or lo[0] >= u.x + u.size or hi[1] < u.y or lo[1] >= u.y + u.size:  # @line L9
                continue
            if u.children is None:  # @line L10
                out.extend(p for p in u.points if lo[0] <= p[0] <= hi[0] and lo[1] <= p[1] <= hi[1])
                continue
            stack.extend(u.children)  # @line L11
        return out

    def delete(self, p: Point) -> bool:
        """Remove p and merge cells that become small; False if p is absent. O(depth * B)."""
        q = self._norm(p)
        u, path = self._leaf(q)  # @line L12
        if q not in u.points:
            return False
        u.points.remove(q)  # @line L13
        u.count -= 1
        for v in path:
            v.count -= 1
        for v in reversed(path):  # @line L14
            if v.count > self.capacity:
                break
            assert v.children is not None
            v.points = [r for c in v.children for r in c.points]
            v.children = None
        return True  # @line L15

    def points(self) -> list[Point]:
        out: list[Point] = []
        stack = [self.root]
        while stack:
            u = stack.pop()
            if u.children is None:
                out.extend(u.points)
            else:
                stack.extend(u.children)
        return out

    def height(self) -> int:
        """Levels below the root (0 when the root is a leaf)."""
        best = 0
        stack = [(self.root, 0)]
        while stack:
            u, d = stack.pop()
            best = max(best, d)
            if u.children is not None:
                stack.extend((c, d + 1) for c in u.children)
        return best

    def node_count(self) -> int:
        n = 0
        stack = [self.root]
        while stack:
            u = stack.pop()
            n += 1
            if u.children is not None:
                stack.extend(u.children)
        return n

    def check_invariant(self) -> None:
        """Raise AssertionError if the structure is broken. O(number of nodes)."""
        seen: set[Point] = set()
        stack = [self.root]
        while stack:
            u = stack.pop()
            if u.children is None:
                assert len(u.points) <= self.capacity and u.count == len(u.points)
                for p in u.points:
                    assert u.holds(p) and p not in seen
                    seen.add(p)
            else:
                assert not u.points and len(u.children) == 4 and u.size >= 2
                h = u.size // 2
                for o, c in enumerate(u.children):
                    assert (c.x, c.y, c.size) == (u.x + (o & 1) * h, u.y + (o >> 1) * h, h)
                    stack.append(c)
                assert u.count == sum(c.count for c in u.children) > self.capacity
        assert len(seen) == self.root.count