Two-dimensional Fenwick tree in Python

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

from collections.abc import Sequence


def _low(i: int) -> int:
    """Return the lowest set bit of i > 0, written g(i) on the page."""
    return i & -i


class FenwickTree2D:
    """A two-dimensional Fenwick tree over an n x m integer matrix.

    The table f is 1-indexed with f[0][*] and f[*][0] unused. The cell f[i][j]
    holds the sum of M over the rows (i - g(i), i] and the columns
    (j - g(j), j], in 1-indexed form, where g(i) is the lowest set bit of i.
    Build takes O(nm) time; add, prefix and query take O(log n log m) time.
    Public methods use 0-indexed rows and columns.
    """

    def __init__(self, matrix: Sequence[Sequence[int]]) -> None:
        self.n = len(matrix)
        self.m = len(matrix[0]) if self.n else 0
        n, m = self.n, self.m
        f = [[0] * (m + 1) for _ in range(n + 1)]
        for i in range(1, n + 1):  # @line L1
            row = matrix[i - 1]
            if len(row) != m:
                raise ValueError("all rows must have the same length")
            for j in range(1, m + 1):
                f[i][j] = row[j - 1]
        for i in range(1, n + 1):  # @line L2
            for j in range(1, m + 1):
                p = j + _low(j)
                if p <= m:  # @line L3
                    f[i][p] += f[i][j]
        for i in range(1, n + 1):  # @line L4
            p = i + _low(i)
            if p <= n:  # @line L5
                fp, fi = f[p], f[i]
                for j in range(1, m + 1):
                    fp[j] += fi[j]
        self.f = f

    def add(self, r: int, c: int, delta: int) -> None:
        """Add delta to M[r][c], 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")
        i = r + 1  # @line L6
        while i <= self.n:  # @line L7
            j = c + 1  # @line L8
            row = self.f[i]
            while j <= self.m:  # @line L9
                row[j] += delta  # @line L10
                j += _low(j)  # @line L11
            i += _low(i)  # @line L12

    def prefix(self, r: int, c: int) -> int:
        """Return the sum of M over the rows 0..r-1 and the columns 0..c-1."""
        if not (0 <= r <= self.n and 0 <= c <= self.m):
            raise IndexError("prefix out of range")
        s = 0  # @line L13
        i = r
        while i > 0:  # @line L14
            j = c  # @line L15
            row = self.f[i]
            while j > 0:  # @line L16
                s += row[j]  # @line L17
                j -= _low(j)  # @line L18
            i -= _low(i)  # @line L19
        return s  # @line L20

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

        An empty rectangle (r1 = r2 + 1 or c1 = c2 + 1) gives 0.
        """
        if not (0 <= r1 <= r2 + 1 <= self.n and 0 <= c1 <= c2 + 1 <= self.m):
            raise IndexError("rectangle out of range")
        return (  # @line L21
            self.prefix(r2 + 1, c2 + 1)
            - self.prefix(r1, c2 + 1)
            - self.prefix(r2 + 1, c1)
            + self.prefix(r1, c1)
        )

    def to_matrix(self) -> list[list[int]]:
        """Return the current matrix M, in O(nm log n log m) time."""
        return [
            [self.query(r, c, r, c) for c in range(self.m)] for r in range(self.n)
        ]

    def check_invariant(self) -> None:
        """Assert that each cell holds the sum of its rectangle of M."""
        mat = self.to_matrix()
        for i in range(1, self.n + 1):
            for j in range(1, self.m + 1):
                want = sum(
                    mat[a][b]
                    for a in range(i - _low(i), i)
                    for b in range(j - _low(j), j)
                )
                assert self.f[i][j] == want, (i, j)