Segment tree beats in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/segment-tree-beats/. The authors are listed in prov/4S5W6M.jsonl.

from collections.abc import Sequence

NEG = float("-inf")


class SegmentTreeBeats:
    """Segment tree beats for range chmin, range sum and range max on integers.

    Node 1 covers A[0 .. n-1]; node v with range [lo, hi] has the children 2v
    and 2v + 1. Each node keeps the maximum mx1, the strict second maximum mx2
    (NEG if all values of the range are equal), the number cnt of items equal
    to mx1, and the sum. A pending chmin at a node is mx1 itself: the true
    value of an item is the minimum of mx1 over the nodes on its root path.
    chmin costs O(log n) amortized; the queries cost O(log n).
    """

    def __init__(self, values: Sequence[int]) -> None:
        self.n = len(values)
        if self.n == 0:
            raise ValueError("the array must not be empty")
        size = 4 * self.n
        self.mx1: list[float] = [NEG] * size
        self.mx2: list[float] = [NEG] * size
        self.cnt = [0] * size
        self.sum = [0] * size
        self.visits = 0  # nodes entered by chmin, for the amortized bound
        self._build(1, 0, self.n - 1, values)

    def _build(self, v: int, lo: int, hi: int, a: Sequence[int]) -> None:
        if lo == hi:  # @line L1
            self.mx1[v], self.mx2[v], self.cnt[v], self.sum[v] = a[lo], NEG, 1, a[lo]
            return
        mid = (lo + hi) // 2  # @line L2
        self._build(2 * v, lo, mid, a)
        self._build(2 * v + 1, mid + 1, hi, a)
        self._pull(v)  # @line L3

    def _pull(self, v: int) -> None:
        a, b = 2 * v, 2 * v + 1  # @line L4
        mx1, mx2, cnt = self.mx1, self.mx2, self.cnt
        self.sum[v] = self.sum[a] + self.sum[b]  # @line L5
        if mx1[a] == mx1[b]:  # @line L6
            mx1[v], cnt[v], mx2[v] = mx1[a], cnt[a] + cnt[b], max(mx2[a], mx2[b])
        elif mx1[a] > mx1[b]:  # @line L7
            mx1[v], cnt[v], mx2[v] = mx1[a], cnt[a], max(mx2[a], mx1[b])
        else:  # @line L8
            mx1[v], cnt[v], mx2[v] = mx1[b], cnt[b], max(mx1[a], mx2[b])

    def _apply(self, v: int, x: int) -> None:
        if x < self.mx1[v]:  # @line L9
            self.sum[v] -= self.cnt[v] * (self.mx1[v] - x)  # type: ignore[assignment]
            self.mx1[v] = x

    def _push(self, v: int) -> None:
        self._apply(2 * v, self.mx1[v])  # type: ignore[arg-type]  # @line L10
        self._apply(2 * v + 1, self.mx1[v])  # type: ignore[arg-type]

    def _chmin(self, v: int, lo: int, hi: int, l: int, r: int, x: int) -> None:
        self.visits += 1
        if r < lo or hi < l or self.mx1[v] <= x:  # @line L11
            return
        if l <= lo and hi <= r and self.mx2[v] < x:  # @line L12
            self._apply(v, x)
            return
        self._push(v)  # @line L13
        mid = (lo + hi) // 2
        self._chmin(2 * v, lo, mid, l, r, x)  # @line L14
        self._chmin(2 * v + 1, mid + 1, hi, l, r, x)
        self._pull(v)  # @line L15

    def chmin(self, l: int, r: int, x: int) -> None:
        """Set A[i] = min(A[i], x) for l <= i <= r."""
        self._check(l, r)
        self._chmin(1, 0, self.n - 1, l, r, x)

    def _query(self, v: int, lo: int, hi: int, l: int, r: int, want_sum: bool) -> float:
        if r < lo or hi < l:  # @line L16
            return 0 if want_sum else NEG
        if l <= lo and hi <= r:  # @line L17
            return self.sum[v] if want_sum else self.mx1[v]
        self._push(v)  # @line L18
        mid = (lo + hi) // 2
        a = self._query(2 * v, lo, mid, l, r, want_sum)
        b = self._query(2 * v + 1, mid + 1, hi, l, r, want_sum)
        return a + b if want_sum else max(a, b)

    def query_sum(self, l: int, r: int) -> int:
        """Return A[l] + ... + A[r]."""
        self._check(l, r)
        return int(self._query(1, 0, self.n - 1, l, r, True))

    def query_max(self, l: int, r: int) -> int:
        """Return max(A[l], ..., A[r])."""
        self._check(l, r)
        return int(self._query(1, 0, self.n - 1, l, r, False))

    def values(self) -> list[int]:
        """Return the current array, in O(n log n) time."""
        return [self.query_max(i, i) for i in range(self.n)]

    def _check(self, l: int, r: int) -> None:
        if not (0 <= l <= r < self.n):
            raise IndexError("range out of bounds")

    def check_invariant(self) -> None:
        """Assert that every node describes the capped values of its range.

        The value of item i seen at node v is the minimum of the leaf value
        and mx1 of the nodes from the leaf up to v. The fields of v must be
        the max, strict second max, count of the max and sum of these values,
        and the values seen at the root must be the array.
        """

        def walk(v: int, lo: int, hi: int) -> list[float]:
            if lo == hi:
                seen = [self.mx1[v]]
            else:
                mid = (lo + hi) // 2
                seen = walk(2 * v, lo, mid) + walk(2 * v + 1, mid + 1, hi)
                seen = [min(s, self.mx1[v]) for s in seen]
            top = max(seen)
            rest = [s for s in seen if s != top]
            assert self.mx1[v] == top, v
            assert self.mx2[v] == (max(rest) if rest else NEG), v
            assert self.cnt[v] == len(seen) - len(rest), v
            assert self.sum[v] == sum(seen), v
            return seen

        assert walk(1, 0, self.n - 1) == self.values()