Segment tree in Python

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

import operator
from collections.abc import Callable, Sequence
from typing import Generic, TypeVar

T = TypeVar("T")


class SegmentTree(Generic[T]):
    """A recursive segment tree over a monoid (op, identity).

    Node 1 covers A[0 .. n-1]; node v with range [lo, hi] has the children
    2v (range [lo, mid]) and 2v + 1 (range [mid + 1, hi]), mid = (lo + hi) // 2.
    t[v] holds the product of A over the range of v. Building takes O(n) time,
    query and update take O(log n) time, and the list has fewer than 4n cells.
    op must be associative and identity must be its neutral element; op need
    not be commutative.
    """

    def __init__(
        self,
        values: Sequence[T],
        op: Callable[[T, T], T] = operator.add,  # type: ignore[assignment]
        identity: T = 0,  # type: ignore[assignment]
    ) -> None:
        self.n = len(values)
        self.op = op
        self.e = identity
        self.a = list(values)
        self.t: list[T] = [identity] * (4 * max(self.n, 1))
        if self.n:
            self._build(1, 0, self.n - 1)

    def _build(self, v: int, lo: int, hi: int) -> None:
        if lo == hi:  # @line L1
            self.t[v] = self.a[lo]
            return
        mid = (lo + hi) // 2  # @line L2
        self._build(2 * v, lo, mid)  # @line L3
        self._build(2 * v + 1, mid + 1, hi)  # @line L4
        self.t[v] = self.op(self.t[2 * v], self.t[2 * v + 1])  # @line L5

    def _query(self, v: int, lo: int, hi: int, l: int, r: int) -> T:
        if r < lo or hi < l:  # @line L6
            return self.e
        if l <= lo and hi <= r:  # @line L7
            return self.t[v]
        mid = (lo + hi) // 2  # @line L8
        return self.op(  # @line L9
            self._query(2 * v, lo, mid, l, r),
            self._query(2 * v + 1, mid + 1, hi, l, r),
        )

    def _update(self, v: int, lo: int, hi: int, i: int, x: T) -> None:
        if lo == hi:  # @line L10
            self.t[v] = x
            return
        mid = (lo + hi) // 2  # @line L11
        if i <= mid:  # @line L12
            self._update(2 * v, lo, mid, i, x)
        else:  # @line L13
            self._update(2 * v + 1, mid + 1, hi, i, x)
        self.t[v] = self.op(self.t[2 * v], self.t[2 * v + 1])  # @line L14

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

    def query(self, l: int, r: int) -> T:
        """Return A[l] op ... op A[r], or the identity when l = r + 1.

        Precondition: 0 <= l <= r + 1 <= n. O(log n) time.
        """
        if not (0 <= l <= r + 1 <= self.n):
            raise IndexError("bad range")
        if l > r:
            return self.e
        return self._query(1, 0, self.n - 1, l, r)

    def update(self, i: int, x: T) -> None:
        """Set A[i] to x. Precondition: 0 <= i < n. O(log n) time."""
        if not 0 <= i < self.n:
            raise IndexError("index out of range")
        self.a[i] = x
        self._update(1, 0, self.n - 1, i, x)

    def check_invariant(self) -> None:
        """Check that each node holds the product of its range."""
        if not self.n:
            return
        stack = [(1, 0, self.n - 1)]
        while stack:
            v, lo, hi = stack.pop()
            want = self.a[lo]
            for k in range(lo + 1, hi + 1):
                want = self.op(want, self.a[k])
            assert self.t[v] == want, (v, lo, hi)
            if lo < hi:
                mid = (lo + hi) // 2
                stack.append((2 * v, lo, mid))
                stack.append((2 * v + 1, mid + 1, hi))