Bitset adjacency matrix in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/bitset-adjacency-matrix/. The authors are listed in prov/WJ5KV3.jsonl.


class BitsetAdjacencyMatrix:
    """A graph on 0, ..., n - 1 whose matrix rows are packed into w-bit words.

    Bit v % w of word rows[u][v // w] is 1 when (u, v) is an edge. adjacent,
    add_edge and remove_edge take O(1) word operations; neighbors takes
    O(n/w + d(u)); degree and common take O(n/w). The space is n * ceil(n/w)
    words.
    """

    def __init__(self, n: int, directed: bool = False, w: int = 64) -> None:
        if n < 0 or w < 1:
            raise ValueError("need n >= 0 and w >= 1")
        self.n = n
        self.w = w
        self.directed = directed
        words = -(-n // w)
        self.rows: list[list[int]] = [[0] * words for _ in range(n)]  # @line L1

    def _check(self, u: int, v: int) -> None:
        if not (0 <= u < self.n and 0 <= v < self.n):
            raise IndexError("vertex out of range")

    def add_edge(self, u: int, v: int) -> None:
        self._check(u, v)
        if not self.directed and u == v:
            raise ValueError("an undirected graph has no self-loops")
        w = self.w
        self.rows[u][v // w] |= 1 << (v % w)  # @line L2
        if not self.directed:
            self.rows[v][u // w] |= 1 << (u % w)  # @line L3

    def remove_edge(self, u: int, v: int) -> None:
        self._check(u, v)
        w = self.w
        self.rows[u][v // w] &= ~(1 << (v % w))  # @line L4
        if not self.directed:
            self.rows[v][u // w] &= ~(1 << (u % w))  # @line L5

    def adjacent(self, u: int, v: int) -> bool:
        self._check(u, v)
        w = self.w
        return (self.rows[u][v // w] >> (v % w)) & 1 == 1  # @line L6

    def neighbors(self, u: int) -> list[int]:
        self._check(u, u)
        out = []
        for j, word in enumerate(self.rows[u]):  # @line L7
            x = word  # @line L8
            while x:  # @line L9
                out.append(j * self.w + (x & -x).bit_length() - 1)  # @line L10
                x &= x - 1  # @line L11
        return out

    def degree(self, u: int) -> int:
        self._check(u, u)
        return sum(x.bit_count() for x in self.rows[u])  # @line L12

    def common(self, u: int, v: int) -> int:
        self._check(u, v)
        return sum((a & b).bit_count() for a, b in zip(self.rows[u], self.rows[v]))  # @line L13

    def add_vertex(self) -> int:
        if self.n % self.w == 0:
            for row in self.rows:
                row.append(0)  # @line L14
        self.rows.append([0] * (-(-(self.n + 1) // self.w)))  # @line L15
        self.n += 1
        return self.n - 1

    def edges(self) -> list[tuple[int, int]]:
        return [
            (u, v)
            for u in range(self.n)
            for v in self.neighbors(u)
            if self.directed or u < v
        ]

    def complement(self) -> "BitsetAdjacencyMatrix":
        """The complement graph: NOT of each word, then clear padding and, if undirected, the diagonal."""
        c = BitsetAdjacencyMatrix(self.n, self.directed, self.w)
        full = (1 << self.w) - 1
        r = self.n % self.w
        last = full if r == 0 else (1 << r) - 1
        for u in range(self.n):
            row = [~x & full for x in self.rows[u]]
            if row:
                row[-1] &= last
            if not self.directed:
                row[u // self.w] &= ~(1 << (u % self.w))
            c.rows[u] = row
        return c

    def check_invariant(self) -> None:
        words = -(-self.n // self.w)
        assert len(self.rows) == self.n
        r = self.n % self.w
        for u in range(self.n):
            row = self.rows[u]
            assert len(row) == words
            for x in row:
                assert 0 <= x < (1 << self.w)
            if row and r:
                assert row[-1] >> r == 0
            if not self.directed:
                assert not self.adjacent(u, u)
                for v in self.neighbors(u):
                    assert self.adjacent(v, u)