Bit vector in Python

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


class BitSet:
    """A set of integers from {0, ..., n - 1} stored as an array of w-bit words.

    Element i is bit (i mod w) of word i // w. Bits at positions n and above
    in the last word are always 0. insert, delete and contains cost O(1);
    count, complement and union cost O(n / w) word operations; next_member
    costs O(n / w) in the worst case. Space: ceil(n / w) words.
    """

    def __init__(self, n: int = 0, 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._mask = (1 << w) - 1
        self.words = [0] * ((n + w - 1) // w)  # @line L1

    def _check_index(self, i: int) -> None:
        if not 0 <= i < self.n:
            raise IndexError("element out of the universe")

    def insert(self, i: int) -> None:
        """Add i to the set. Precondition: 0 <= i < n. Cost O(1)."""
        self._check_index(i)
        self.words[i // self.w] |= 1 << (i % self.w)  # @line L2

    def delete(self, i: int) -> None:
        """Remove i from the set if present. Precondition: 0 <= i < n. Cost O(1)."""
        self._check_index(i)
        self.words[i // self.w] &= ~(1 << (i % self.w)) & self._mask  # @line L3

    def contains(self, i: int) -> bool:
        """Return True if i is in the set. Precondition: 0 <= i < n. Cost O(1)."""
        self._check_index(i)
        return (self.words[i // self.w] >> (i % self.w)) & 1 == 1  # @line L4

    def count(self) -> int:
        """Return the number of members. Cost O(n / w) popcounts."""
        c = 0
        for word in self.words:
            c += word.bit_count()
        return c

    def next_member(self, i: int) -> int | None:
        """Return the least member >= i, or None. Precondition: i >= 0. Cost O(n / w)."""
        j = i // self.w
        if j >= len(self.words):
            return None
        x = self.words[j] & ~((1 << (i % self.w)) - 1) & self._mask
        while x == 0:
            j += 1
            if j == len(self.words):
                return None
            x = self.words[j]
        return j * self.w + ((x & -x).bit_length() - 1)

    def complement(self) -> None:
        """Replace the set by its complement in {0, ..., n - 1}. Cost O(n / w)."""
        for j in range(len(self.words)):
            self.words[j] = ~self.words[j] & self._mask
        r = self.n % self.w
        if r != 0:
            self.words[-1] &= (1 << r) - 1

    def union(self, other: "BitSet") -> None:
        """Add every member of other, a set over the same n and w. Cost O(n / w)."""
        if other.n != self.n or other.w != self.w:
            raise ValueError("universes differ")
        for j in range(len(self.words)):
            self.words[j] |= other.words[j]

    def members(self) -> list[int]:
        """Return the members in increasing order. Cost O(n / w + size)."""
        out = []
        for j, word in enumerate(self.words):
            while word:
                b = word & -word
                out.append(j * self.w + b.bit_length() - 1)
                word ^= b
        return out

    def check_invariant(self) -> None:
        """Assert the representation invariant: the right number of words,
        each a w-bit word, and no bit at a position >= n."""
        assert len(self.words) == (self.n + self.w - 1) // self.w
        for word in self.words:
            assert 0 <= word <= self._mask
        r = self.n % self.w
        if r != 0 and self.words:
            assert self.words[-1] >> r == 0