Dynamic (sparse) segment tree in Python

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


class DynamicSegmentTree:
    """Range sums over A[0..U-1] of integers, all zero at the start.

    Nodes are made only on the paths of updates. update and query take
    O(log U) time; after m updates the tree has at most 1 + m*ceil(lg U)
    nodes. Node 0 means "absent" and node 1 is the root.
    """

    def __init__(self, universe: int) -> None:
        """Make the tree of U = universe zeros; requires universe >= 1."""
        if universe < 1:
            raise ValueError("universe must be at least 1")
        self.universe = universe
        self.val = [0, 0]
        self.left = [0, 0]
        self.right = [0, 0]

    def _new_node(self) -> int:
        self.val.append(0)
        self.left.append(0)
        self.right.append(0)
        return len(self.val) - 1

    def node_count(self) -> int:
        """The number of created nodes, the root included."""
        return len(self.val) - 1

    def update(self, i: int, x: int) -> None:
        """Set A[i] += x, for 0 <= i < U, in O(log U) time."""
        if not 0 <= i < self.universe:
            raise IndexError("index out of range")
        v, lo, hi = 1, 0, self.universe  # @line L1
        while True:  # @line L2
            self.val[v] += x  # @line L3
            if hi - lo == 1:  # @line L4
                return
            mid = (lo + hi) // 2  # @line L5
            if i < mid:  # @line L6
                if self.left[v] == 0:  # @line L7
                    c = self._new_node()
                    self.left[v] = c
                v, hi = self.left[v], mid  # @line L8
            else:  # @line L9
                if self.right[v] == 0:  # @line L10
                    c = self._new_node()
                    self.right[v] = c
                v, lo = self.right[v], mid  # @line L11

    def query(self, l: int, r: int) -> int:
        """Return A[l] + ... + A[r-1] for 0 <= l <= r <= U, in O(log U) time."""
        if not 0 <= l <= r <= self.universe:
            raise IndexError("range out of bounds")
        total = 0
        stack = [(1, 0, self.universe)]
        while stack:
            v, lo, hi = stack.pop()
            if v == 0 or r <= lo or hi <= l:  # @line L12
                continue
            if l <= lo and hi <= r:  # @line L13
                total += self.val[v]
                continue
            mid = (lo + hi) // 2  # @line L14
            stack.append((self.right[v], mid, hi))  # @line L15
            stack.append((self.left[v], lo, mid))
        return total

    def check_invariant(self) -> None:
        """Check that each inner node holds the sum of its children and
        that each node is reached exactly once from the root."""
        seen = [False] * len(self.val)
        stack = [(1, 0, self.universe)]
        while stack:
            v, lo, hi = stack.pop()
            assert not seen[v], "node reached twice"
            seen[v] = True
            if hi - lo == 1:
                assert self.left[v] == 0 and self.right[v] == 0, "leaf with children"
                continue
            mid = (lo + hi) // 2
            s = 0
            for c, a, b in ((self.left[v], lo, mid), (self.right[v], mid, hi)):
                if c != 0:
                    s += self.val[c]
                    stack.append((c, a, b))
            if self.left[v] or self.right[v]:
                assert self.val[v] == s, "inner sum differs from the children"
        assert all(seen[1:]), "unreachable node"