Dynamic bit vector in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/dynamic-bit-vector/. The authors are listed in prov/ZZJ07J.jsonl.
"""Dynamic bit vector: rank, select, insert and delete in O(log n) expected time.

The bits are cut into chunks of at most 2B bits. The chunks, in order,
are the nodes of an implicit treap; each node also stores the number of
bits and of 1s in its subtree. A query walks down from the root and
finishes inside one chunk; an update changes one chunk, fixes the counts
on the path, and splits a chunk that grows past 2B bits.
"""

from __future__ import annotations

import random
from typing import Optional


class _Node:
    __slots__ = ("chunk", "pri", "left", "right", "size", "ones")

    def __init__(self, chunk: list[int], pri: float) -> None:
        self.chunk = chunk
        self.pri = pri
        self.left: Optional[_Node] = None
        self.right: Optional[_Node] = None
        self.size = len(chunk)
        self.ones = sum(chunk)


def _size(t: Optional[_Node]) -> int:
    return t.size if t else 0


def _ones(t: Optional[_Node]) -> int:
    return t.ones if t else 0


def _chunk_ones(t: _Node) -> int:
    """Number of 1s in the node's own chunk."""
    return t.ones - _ones(t.left) - _ones(t.right)


def _pull(t: _Node) -> None:
    t.size = _size(t.left) + len(t.chunk) + _size(t.right)
    t.ones = _ones(t.left) + sum(t.chunk) + _ones(t.right)


def _split(t: Optional[_Node], k: int) -> tuple[Optional[_Node], Optional[_Node]]:
    """Split at bit k, which must be a chunk boundary: left has bits [0, k)."""
    if t is None:
        return None, None
    ls = _size(t.left)
    if k <= ls:
        a, b = _split(t.left, k)
        t.left = b
        _pull(t)
        return a, t
    a, b = _split(t.right, k - ls - len(t.chunk))
    t.right = a
    _pull(t)
    return t, b


def _merge(a: Optional[_Node], b: Optional[_Node]) -> Optional[_Node]:
    if a is None:
        return b
    if b is None:
        return a
    if a.pri > b.pri:
        a.right = _merge(a.right, b)
        _pull(a)
        return a
    b.left = _merge(a, b.left)
    _pull(b)
    return b


def _pop_leftmost(t: _Node) -> Optional[_Node]:
    if t.left is None:
        return t.right
    t.left = _pop_leftmost(t.left)
    _pull(t)
    return t


class DynamicBitVector:
    """A bit vector with access, rank, select, insert, delete and set."""

    def __init__(self, bits: Optional[list[int]] = None, block: int = 64, seed: int = 0) -> None:
        """Build from ``bits``; chunks hold between 1 and 2 * ``block`` bits."""
        self.B = block
        self.rng = random.Random(seed)
        self.root: Optional[_Node] = None
        bits = list(bits or [])
        for s in range(0, len(bits), block):
            self.root = _merge(self.root, _Node(bits[s:s + block], self.rng.random()))

    def __len__(self) -> int:
        return _size(self.root)

    def to_list(self) -> list[int]:
        """All bits, in order."""
        out: list[int] = []

        def walk(t: Optional[_Node]) -> None:
            if t:
                walk(t.left)
                out.extend(t.chunk)
                walk(t.right)

        walk(self.root)
        return out

    def _find(self, i: int) -> tuple[list[_Node], int, int]:
        """Path to the node holding bit i (or the last node if i = n), offset, start."""
        path: list[_Node] = []
        x, start = self.root, 0
        while x is not None:
            path.append(x)
            ls = _size(x.left)
            if i < ls:
                x = x.left
            elif i < ls + len(x.chunk) or (x.right is None and i == ls + len(x.chunk)):
                return path, i - ls, start + ls
            else:
                i -= ls + len(x.chunk)
                start += ls + len(x.chunk)
                x = x.right
        raise IndexError("position out of range")

    def access(self, i: int) -> int:
        """The bit at position i."""
        if not 0 <= i < len(self):
            raise IndexError("position out of range")
        path, off, _ = self._find(i)
        return path[-1].chunk[off]

    def rank1(self, i: int) -> int:
        """Number of 1s in positions 0..i-1, for 0 <= i <= n."""
        x, r = self.root, 0  # @line L1
        while x is not None:  # @line L2
            ls = _size(x.left)
            if i < ls:  # @line L3
                x = x.left
            elif i <= ls + len(x.chunk):  # @line L4
                return r + _ones(x.left) + sum(x.chunk[: i - ls])
            else:  # @line L5
                r += _ones(x.left) + _chunk_ones(x)
                i -= ls + len(x.chunk)
                x = x.right
        return r  # @line L6

    def rank0(self, i: int) -> int:
        """Number of 0s in positions 0..i-1."""
        return i - self.rank1(i)

    def select1(self, j: int) -> int:
        """Position of the j-th 1, for 1 <= j <= rank1(n)."""
        if not 1 <= j <= _ones(self.root):
            raise IndexError("no such 1")
        x, base = self.root, 0  # @line L7
        while x is not None:  # @line L8
            lo, ls = _ones(x.left), _size(x.left)
            co = _chunk_ones(x)
            if j <= lo:  # @line L9
                x = x.left
            elif j <= lo + co:  # @line L10
                j -= lo
                for k, b in enumerate(x.chunk):
                    j -= b
                    if j == 0:
                        return base + ls + k
            else:  # @line L11
                j -= lo + co
                base += ls + len(x.chunk)
                x = x.right
        raise AssertionError("unreachable")

    def select0(self, j: int) -> int:
        """Position of the j-th 0, by binary search over rank0."""
        if not 1 <= j <= len(self) - _ones(self.root):
            raise IndexError("no such 0")
        lo, hi = 0, len(self) - 1
        while lo < hi:
            mid = (lo + hi) // 2
            if self.rank0(mid + 1) >= j:
                hi = mid
            else:
                lo = mid + 1
        return lo

    def insert(self, i: int, b: int) -> None:
        """Insert bit b before position i, for 0 <= i <= n."""
        if not 0 <= i <= len(self):
            raise IndexError("position out of range")
        if self.root is None:
            self.root = _Node([b], self.rng.random())
            return
        path, off, start = self._find(i)  # @line L12
        x = path[-1]
        x.chunk.insert(off, b)  # @line L13
        for y in reversed(path):  # @line L14
            _pull(y)
        if len(x.chunk) > 2 * self.B:  # @line L15
            half = len(x.chunk) // 2
            new = _Node(x.chunk[half:], self.rng.random())
            x.chunk = x.chunk[:half]
            for y in reversed(path):
                _pull(y)
            left, right = _split(self.root, start + half)
            self.root = _merge(_merge(left, new), right)

    def delete(self, i: int) -> int:
        """Remove and return the bit at position i."""
        if not 0 <= i < len(self):
            raise IndexError("position out of range")
        path, off, start = self._find(i)  # @line L16
        x = path[-1]
        b = x.chunk.pop(off)  # @line L17
        for y in reversed(path):  # @line L18
            _pull(y)
        if not x.chunk:  # @line L19
            left, right = _split(self.root, start)
            assert right is not None
            self.root = _merge(left, _pop_leftmost(right))
        return b

    def set(self, i: int, b: int) -> None:
        """Overwrite the bit at position i with b."""
        if not 0 <= i < len(self):
            raise IndexError("position out of range")
        path, off, _ = self._find(i)
        path[-1].chunk[off] = b
        for y in reversed(path):
            _pull(y)