Linear quadtree in Python

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

from bisect import bisect_left, bisect_right
from typing import Iterable

Pixel = tuple[int, int]
Block = tuple[int, int]  # (Morton code of the lower left pixel, level)


def morton(x: int, y: int, depth: int) -> int:
    """Interleave the bits of x and y: bit i of x goes to bit 2i, bit i of y to bit 2i+1."""
    z = 0
    for i in range(depth):
        z |= ((x >> i) & 1) << (2 * i)
        z |= ((y >> i) & 1) << (2 * i + 1)
    return z


def demorton(z: int, depth: int) -> Pixel:
    """The inverse of morton."""
    x = y = 0
    for i in range(depth):
        x |= ((z >> (2 * i)) & 1) << i
        y |= ((z >> (2 * i + 1)) & 1) << i
    return x, y


class LinearQuadtree:
    """The black pixels of a 2^depth by 2^depth binary image, as a sorted list of blocks.

    A block (c, l) is the aligned square of side 2^l whose pixels have the Morton
    codes c to c + 4^l - 1. Invariant: c is a multiple of 4^l, the blocks are
    disjoint and sorted by code, and no four sibling blocks of one level are all
    present (they would form their parent). So the list is the set of black
    leaves of the region quadtree of the image, and it is unique.
    """

    def __init__(self, depth: int) -> None:
        if depth < 0:
            raise ValueError("need depth >= 0")
        self.depth = depth
        self.codes: list[int] = []
        self.levels: list[int] = []

    @classmethod
    def from_pixels(cls, depth: int, pixels: Iterable[Pixel]) -> "LinearQuadtree":
        """Build from a set of black pixels with one sort and one stack pass. O(n log n)."""
        t = cls(depth)
        zs = sorted({morton(*t._norm(p), depth) for p in pixels})
        stack: list[Block] = []
        for z in zs:  # @line L1
            stack.append((z, 0))  # @line L2
            while len(stack) >= 4:  # @line L3
                c, l = stack[-4]
                step = 1 << (2 * l)
                if c % (step << 2) != 0 or any(stack[-4 + k] != (c + k * step, l) for k in range(1, 4)):  # @line L4
                    break
                del stack[-4:]  # @line L5
                stack.append((c, l + 1))  # @line L6
        t.codes = [c for c, _ in stack]
        t.levels = [l for _, l in stack]
        return t

    def __len__(self) -> int:
        return len(self.codes)

    def _norm(self, p: Pixel) -> Pixel:
        x, y = int(p[0]), int(p[1])
        if not (0 <= x < (1 << self.depth) and 0 <= y < (1 << self.depth)):
            raise ValueError("the pixel is outside the image")
        return x, y

    def _locate(self, z: int) -> int:
        """Index of the block that holds the code z, or -1."""
        i = bisect_right(self.codes, z) - 1  # @line L7
        if i >= 0 and z < self.codes[i] + (1 << (2 * self.levels[i])):  # @line L8
            return i
        return -1  # @line L9

    def locate(self, p: Pixel) -> tuple[int, int, int] | None:
        """The block (x, y, side) that holds the pixel p, or None if p is white. O(log n)."""
        i = self._locate(morton(*self._norm(p), self.depth))
        if i < 0:
            return None
        x, y = demorton(self.codes[i], self.depth)
        return x, y, 1 << self.levels[i]

    def __contains__(self, p: Pixel) -> bool:
        return self._locate(morton(*self._norm(p), self.depth)) >= 0

    def _has(self, c: int, l: int) -> bool:
        i = bisect_left(self.codes, c)
        return i < len(self.codes) and self.codes[i] == c and self.levels[i] == l

    def _add(self, c: int, l: int) -> None:
        i = bisect_left(self.codes, c)
        self.codes.insert(i, c)
        self.levels.insert(i, l)

    def _remove(self, c: int) -> None:
        i = bisect_left(self.codes, c)
        del self.codes[i]
        del self.levels[i]

    def insert(self, p: Pixel) -> bool:
        """Make the pixel p black. Returns False if it was black. O(depth) list updates."""
        z = morton(*self._norm(p), self.depth)
        if self._locate(z) >= 0:  # @line L10
            return False
        c, l = z, 0
        self._add(c, l)  # @line L11
        while l < self.depth:
            step = 1 << (2 * l)
            parent = c - c % (step << 2)
            if not all(self._has(parent + k * step, l) for k in range(4)):  # @line L12
                break
            for k in range(4):  # @line L13
                self._remove(parent + k * step)
            c, l = parent, l + 1
            self._add(c, l)
        return True

    def delete(self, p: Pixel) -> bool:
        """Make the pixel p white. Returns False if it was white. O(depth) list updates."""
        z = morton(*self._norm(p), self.depth)
        i = self._locate(z)  # @line L14
        if i < 0:
            return False
        l = self.levels[i]
        self._remove(self.codes[i])  # @line L15
        for j in range(l - 1, -1, -1):  # @line L16
            step = 1 << (2 * j)
            base = z - z % (step << 2)
            for k in range(4):  # @line L17
                if base + k * step != z - z % step:
                    self._add(base + k * step, j)
        return True

    def blocks_in_cell(self, x: int, y: int, level: int) -> list[tuple[int, int, int]]:
        """The black blocks inside the aligned cell of side 2^level at (x, y), clipped to it.

        The blocks inside a cell form one run of the sorted list, found by two
        binary searches. A block that covers the whole cell is returned as the cell.
        """
        if level < 0 or level > self.depth or x % (1 << level) or y % (1 << level):
            raise ValueError("need an aligned cell inside the image")
        lo = morton(*self._norm((x, y)), self.depth)
        hi = lo + (1 << (2 * level))
        i = self._locate(lo)
        if i >= 0 and self.levels[i] >= level:
            return [(x, y, 1 << level)]
        out = []
        for j in range(bisect_left(self.codes, lo), bisect_left(self.codes, hi)):
            bx, by = demorton(self.codes[j], self.depth)
            out.append((bx, by, 1 << self.levels[j]))
        return out

    def blocks(self) -> list[tuple[int, int, int]]:
        """All blocks as (x, y, side), in Morton order."""
        return [(*demorton(c, self.depth), 1 << l) for c, l in zip(self.codes, self.levels)]

    def area(self) -> int:
        """The number of black pixels."""
        return sum(1 << (2 * l) for l in self.levels)

    def pixels(self) -> list[Pixel]:
        out = []
        for c, l in zip(self.codes, self.levels):
            for z in range(c, c + (1 << (2 * l))):
                out.append(demorton(z, self.depth))
        return out

    def check_invariant(self) -> None:
        full = 1 << (2 * self.depth)
        assert len(self.codes) == len(self.levels)
        end = 0
        for c, l in zip(self.codes, self.levels):
            assert 0 <= l <= self.depth, "level out of range"
            assert c % (1 << (2 * l)) == 0, "block not aligned"
            assert c >= end, "blocks overlap or are out of order"
            end = c + (1 << (2 * l))
            assert end <= full, "block outside the image"
        for c, l in zip(self.codes, self.levels):
            if l < self.depth:
                step = 1 << (2 * l)
                parent = c - c % (step << 2)
                assert not all(self._has(parent + k * step, l) for k in range(4)), "four siblings not merged"