Chtholly tree in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/chtholly-tree/. The authors are listed in prov/MQYKYZ.jsonl.

from bisect import bisect_right
from typing import Callable, Sequence


class ChthollyTree:
    """Runs (start, value) of an integer array A[0..n-1]; run i covers
    [starts[i], starts[i+1]) with starts[R] read as n.

    The runs live in two sorted Python lists, so a split is a binary
    search plus a list insert. assign deletes the runs it covers in one
    slice. Range operations cost O(k + log n) searches for k runs, plus
    the list shifts of Python.
    """

    def __init__(self, a: Sequence[int]) -> None:
        """Make one run per maximal block of equal values of a; n >= 1."""
        if len(a) == 0:
            raise ValueError("the array must be non-empty")
        self.n = len(a)
        self.starts: list[int] = []
        self.vals: list[int] = []
        for i, x in enumerate(a):
            if i == 0 or x != a[i - 1]:
                self.starts.append(i)
                self.vals.append(x)

    def _split(self, x: int) -> int:
        """Make a run start at x, 0 <= x <= n; return its index (R if x = n)."""
        if x == self.n:  # @line L1
            return len(self.starts)
        i = bisect_right(self.starts, x) - 1  # @line L2
        if self.starts[i] == x:  # @line L3
            return i
        self.starts.insert(i + 1, x)  # @line L4
        self.vals.insert(i + 1, self.vals[i])
        return i + 1  # @line L5

    def _check_range(self, l: int, r: int) -> None:
        if not 0 <= l < r <= self.n:
            raise IndexError("need 0 <= l < r <= n")

    def assign(self, l: int, r: int, v: int) -> None:
        """Set A[k] = v for l <= k < r. O(log n) amortized searches."""
        self._check_range(l, r)
        hi = self._split(r)  # @line L6
        before = len(self.starts)
        lo = self._split(l)  # @line L7
        hi += len(self.starts) - before  # an index shift, as l < r
        del self.starts[lo:hi]  # @line L8
        del self.vals[lo:hi]
        self.starts.insert(lo, l)  # @line L9
        self.vals.insert(lo, v)

    def each(self, l: int, r: int, f: Callable[[int, int, int], int]) -> None:
        """Call v = f(s, e, v) on each run (s, e, v) that tiles [l, r)."""
        self._check_range(l, r)
        self._split(r)  # @line L10
        lo = self._split(l)  # @line L11
        i = lo
        while i < len(self.starts) and self.starts[i] < r:  # @line L12
            e = self.starts[i + 1] if i + 1 < len(self.starts) else self.n
            self.vals[i] = f(self.starts[i], e, self.vals[i])  # @line L13
            i += 1

    def add(self, l: int, r: int, d: int) -> None:
        """Add d to A[k] for l <= k < r."""
        self.each(l, r, lambda s, e, v: v + d)

    def range_sum(self, l: int, r: int) -> int:
        """Return A[l] + ... + A[r-1]."""
        total = 0

        def f(s: int, e: int, v: int) -> int:
            nonlocal total
            total += v * (e - s)
            return v

        self.each(l, r, f)
        return total

    def kth_smallest(self, l: int, r: int, k: int) -> int:
        """The k-th smallest of A[l..r-1], 1 <= k <= r - l."""
        if not 1 <= k <= r - l:
            raise ValueError("k out of range")
        parts: list[tuple[int, int]] = []
        self.each(l, r, lambda s, e, v: (parts.append((v, e - s)), v)[1])
        parts.sort()
        for v, c in parts:
            if k <= c:
                return v
            k -= c
        raise AssertionError("unreachable")

    def to_list(self) -> list[int]:
        """The array A."""
        out: list[int] = []
        for i, s in enumerate(self.starts):
            e = self.starts[i + 1] if i + 1 < len(self.starts) else self.n
            out.extend([self.vals[i]] * (e - s))
        return out

    def runs(self) -> list[tuple[int, int, int]]:
        """The runs (s, e, v) in order."""
        ends = self.starts[1:] + [self.n]
        return list(zip(self.starts, ends, self.vals))

    def check_invariant(self) -> None:
        """Check (C): starts strictly increase from 0 and stay below n."""
        assert len(self.starts) == len(self.vals) >= 1
        assert self.starts[0] == 0
        for i in range(1, len(self.starts)):
            assert self.starts[i - 1] < self.starts[i] < self.n