Bit vector in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/bit-vector/. The authors are listed in prov/6G48XR.jsonl.

from collections.abc import Iterable

W = 64  # bits in a word
MASK = (1 << W) - 1


class BitVector:
    """A static bit vector packed into 64-bit words, with access, rank and select.

    Bit i is bit i mod 64 of word i // 64, least significant bit first.
    Bits past the end of the last word are 0.
    Cost: build O(n), access O(1), rank and select O(n / 64 + 64).
    """

    def __init__(self, bits: Iterable[int] = ()) -> None:
        b = list(bits)
        self.n = len(b)
        self.words = [0] * ((self.n + W - 1) // W)  # @line L1
        for i in range(self.n):  # @line L2
            if b[i]:  # @line L3
                self.words[i // W] |= 1 << (i % W)
        self.ones = sum(x.bit_count() for x in self.words)

    def __len__(self) -> int:
        return self.n

    def access(self, i: int) -> int:
        """Return B[i] for 0 <= i < n."""
        if not 0 <= i < self.n:
            raise IndexError(i)
        return (self.words[i // W] >> (i % W)) & 1  # @line L4

    def rank1(self, i: int) -> int:
        """Return the number of 1s in B[0 .. i-1], for 0 <= i <= n."""
        if not 0 <= i <= self.n:
            raise IndexError(i)
        q, r, c = i // W, i % W, 0  # @line L5
        for k in range(q):  # @line L6
            c += self.words[k].bit_count()
        if r > 0:  # @line L7
            c += (self.words[q] & ((1 << r) - 1)).bit_count()
        return c  # @line L8

    def rank0(self, i: int) -> int:
        """Return the number of 0s in B[0 .. i-1], for 0 <= i <= n."""
        return i - self.rank1(i)

    def _word(self, k: int, b: int) -> int:
        x = self.words[k]
        if b == 1:
            return x
        x = ~x & MASK
        if k == len(self.words) - 1 and self.n % W:
            x &= (1 << (self.n % W)) - 1  # mask the padding
        return x

    def _select(self, j: int, b: int) -> int:
        total = self.ones if b == 1 else self.n - self.ones
        if not 1 <= j <= total:
            raise IndexError(j)
        k = 0  # @line L9
        while self._word(k, b).bit_count() < j:  # @line L10
            j -= self._word(k, b).bit_count()
            k += 1
        x = self._word(k, b)  # @line L11
        for _ in range(j - 1):  # @line L12
            x &= x - 1
        return k * W + ((x & -x).bit_length() - 1)  # @line L13

    def select1(self, j: int) -> int:
        """Return the position of the j-th 1, for 1 <= j <= rank1(n)."""
        return self._select(j, 1)

    def select0(self, j: int) -> int:
        """Return the position of the j-th 0, for 1 <= j <= rank0(n)."""
        return self._select(j, 0)

    def check_invariant(self) -> None:
        assert len(self.words) == (self.n + W - 1) // W
        if self.words and self.n % W:
            assert self.words[-1] >> (self.n % W) == 0
        assert all(0 <= x <= MASK for x in self.words)