# 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