Two-dimensional segment tree in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/two-dimensional-segment-tree/. The authors are listed in prov/QYYQT1.jsonl.

import operator
from collections.abc import Callable, Sequence
from typing import Generic, TypeVar

T = TypeVar("T")


class SegmentTree2D(Generic[T]):
    """A segment tree over rows whose nodes are segment trees over columns.

    Outer node vx covers a range of rows and inner node vy a range of
    columns, both numbered as in a binary heap (root 1, children 2v, 2v + 1).
    t[vx][vy] holds the op-product of M over the rows of vx and the columns of
    vy. op must be associative and commutative with the neutral element
    identity. Build takes O(nm) time; query and point assignment take
    O(log n log m) time.
    """

    def __init__(
        self,
        matrix: Sequence[Sequence[T]],
        op: Callable[[T, T], T] = operator.add,  # type: ignore[assignment]
        identity: T = 0,  # type: ignore[assignment]
    ) -> None:
        self.n = len(matrix)
        self.m = len(matrix[0]) if self.n else 0
        if self.n == 0 or self.m == 0:
            raise ValueError("the matrix must have at least one row and one column")
        for row in matrix:
            if len(row) != self.m:
                raise ValueError("all rows must have the same length")
        self.op = op
        self.e = identity
        self.mat = [list(row) for row in matrix]
        self.t: list[list[T]] = [
            [identity] * (4 * self.m) for _ in range(4 * self.n)
        ]
        self._build_x(1, 0, self.n - 1)

    # ----- build -----

    def _build_y(self, vx: int, lx: int, rx: int, vy: int, ly: int, ry: int) -> None:
        t = self.t
        if ly == ry:  # @line L1
            if lx == rx:  # @line L2
                t[vx][vy] = self.mat[lx][ly]
            else:
                t[vx][vy] = self.op(t[2 * vx][vy], t[2 * vx + 1][vy])
            return  # @line L3
        my = (ly + ry) // 2  # @line L4
        self._build_y(vx, lx, rx, 2 * vy, ly, my)
        self._build_y(vx, lx, rx, 2 * vy + 1, my + 1, ry)
        t[vx][vy] = self.op(t[vx][2 * vy], t[vx][2 * vy + 1])  # @line L5

    def _build_x(self, vx: int, lx: int, rx: int) -> None:
        if lx != rx:  # @line L6
            mx = (lx + rx) // 2
            self._build_x(2 * vx, lx, mx)
            self._build_x(2 * vx + 1, mx + 1, rx)
        self._build_y(vx, lx, rx, 1, 0, self.m - 1)  # @line L7

    # ----- query -----

    def _query_y(self, vx: int, vy: int, ly: int, ry: int, c1: int, c2: int) -> T:
        if c2 < ly or ry < c1:  # @line L8
            return self.e
        if c1 <= ly and ry <= c2:  # @line L9
            return self.t[vx][vy]
        my = (ly + ry) // 2  # @line L10
        return self.op(
            self._query_y(vx, 2 * vy, ly, my, c1, c2),
            self._query_y(vx, 2 * vy + 1, my + 1, ry, c1, c2),
        )

    def _query_x(
        self, vx: int, lx: int, rx: int, r1: int, r2: int, c1: int, c2: int
    ) -> T:
        if r2 < lx or rx < r1:  # @line L11
            return self.e
        if r1 <= lx and rx <= r2:  # @line L12
            return self._query_y(vx, 1, 0, self.m - 1, c1, c2)
        mx = (lx + rx) // 2  # @line L13
        return self.op(
            self._query_x(2 * vx, lx, mx, r1, r2, c1, c2),
            self._query_x(2 * vx + 1, mx + 1, rx, r1, r2, c1, c2),
        )

    def query(self, r1: int, c1: int, r2: int, c2: int) -> T:
        """Return the op-product of M[r][c] for r1 <= r <= r2, c1 <= c <= c2.

        An empty rectangle (r1 = r2 + 1 or c1 = c2 + 1) gives the identity.
        """
        if not (0 <= r1 <= r2 + 1 <= self.n and 0 <= c1 <= c2 + 1 <= self.m):
            raise IndexError("rectangle out of range")
        if r1 > r2 or c1 > c2:
            return self.e
        return self._query_x(1, 0, self.n - 1, r1, r2, c1, c2)

    # ----- update -----

    def _update_y(
        self, vx: int, lx: int, rx: int, vy: int, ly: int, ry: int, c: int
    ) -> None:
        t = self.t
        if ly == ry:  # @line L14
            if lx == rx:  # @line L15
                t[vx][vy] = self.mat[lx][ly]
            else:
                t[vx][vy] = self.op(t[2 * vx][vy], t[2 * vx + 1][vy])
            return  # @line L16
        my = (ly + ry) // 2  # @line L17
        if c <= my:
            self._update_y(vx, lx, rx, 2 * vy, ly, my, c)
        else:
            self._update_y(vx, lx, rx, 2 * vy + 1, my + 1, ry, c)
        t[vx][vy] = self.op(t[vx][2 * vy], t[vx][2 * vy + 1])  # @line L18

    def _update_x(self, vx: int, lx: int, rx: int, r: int, c: int) -> None:
        if lx != rx:  # @line L19
            mx = (lx + rx) // 2
            if r <= mx:
                self._update_x(2 * vx, lx, mx, r, c)
            else:
                self._update_x(2 * vx + 1, mx + 1, rx, r, c)
        self._update_y(vx, lx, rx, 1, 0, self.m - 1, c)  # @line L20

    def update(self, r: int, c: int, x: T) -> None:
        """Set M[r][c] = x, for 0 <= r < n and 0 <= c < m."""
        if not (0 <= r < self.n and 0 <= c < self.m):
            raise IndexError("cell out of range")
        self.mat[r][c] = x
        self._update_x(1, 0, self.n - 1, r, c)

    # ----- checks -----

    def _fold(self, r1: int, r2: int, c1: int, c2: int) -> T:
        acc = self.e
        for r in range(r1, r2 + 1):
            for c in range(c1, c2 + 1):
                acc = self.op(acc, self.mat[r][c])
        return acc

    def check_invariant(self) -> None:
        """Assert that each used cell holds the product of its rectangle."""

        def rows(vx: int, lx: int, rx: int) -> None:
            cols(vx, lx, rx, 1, 0, self.m - 1)
            if lx != rx:
                mx = (lx + rx) // 2
                rows(2 * vx, lx, mx)
                rows(2 * vx + 1, mx + 1, rx)

        def cols(vx: int, lx: int, rx: int, vy: int, ly: int, ry: int) -> None:
            assert self.t[vx][vy] == self._fold(lx, rx, ly, ry), (vx, vy)
            if ly != ry:
                my = (ly + ry) // 2
                cols(vx, lx, rx, 2 * vy, ly, my)
                cols(vx, lx, rx, 2 * vy + 1, my + 1, ry)

        rows(1, 0, self.n - 1)