Forward star in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/forward-star/. The authors are listed in prov/83VY7T.jsonl.

from collections.abc import Iterable


def _prefix_counts(n: int, keys: Iterable[int]) -> list[int]:
    """p[i] = the number of keys less than i, for i in 0 .. n."""
    p = [0] * (n + 1)
    for k in keys:
        p[k + 1] += 1
    for i in range(1, n + 1):
        p[i] += p[i - 1]
    return p


class ForwardStar:
    """A static directed network: arcs sorted by tail, named 0 .. m - 1.

    point[i] is the name of the first arc out of i; trace[rpoint[j]:rpoint[j + 1]]
    lists the names of the arcs into j. Building takes Theta(n + m) time,
    out_arcs and in_arcs take time linear in their output, degrees O(1).
    Parallel arcs and self-loops are allowed.
    """

    def __init__(self, n: int, arcs: Iterable[tuple[int, int, float]]) -> None:
        if n < 0:
            raise ValueError("n must be at least 0")
        given = list(arcs)
        for t, h, _ in given:
            if not (0 <= t < n and 0 <= h < n):
                raise IndexError("vertex out of range")
        self.n, self.m = n, len(given)
        start = _prefix_counts(n, (t for t, _, _ in given))
        pos = start[:n]
        order: list[tuple[int, int, float]] = [(0, 0, 0)] * self.m
        for a in given:  # stable counting sort by tail
            order[pos[a[0]]] = a  # @line L1
            pos[a[0]] += 1
        self.tail = [t for t, _, _ in order]
        self.head = [h for _, h, _ in order]
        self.cost = [c for _, _, c in order]
        self.flow: list[float] = [0] * self.m
        self.point = start  # @line L2
        self.rpoint = _prefix_counts(n, self.head)  # @line L3
        rpos = self.rpoint[:n]  # @line L4
        self.trace = [0] * self.m
        for e in range(self.m):  # @line L5
            j = self.head[e]  # @line L6
            self.trace[rpos[j]] = e
            rpos[j] += 1

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

    def out_arcs(self, i: int) -> list[int]:
        self._check(i)
        return list(range(self.point[i], self.point[i + 1]))  # @line L7

    def in_arcs(self, j: int) -> list[int]:
        self._check(j)
        return [self.trace[k] for k in range(self.rpoint[j], self.rpoint[j + 1])]  # @line L8

    def out_degree(self, i: int) -> int:
        self._check(i)
        return self.point[i + 1] - self.point[i]  # @line L9

    def in_degree(self, j: int) -> int:
        self._check(j)
        return self.rpoint[j + 1] - self.rpoint[j]  # @line L10

    def arcs(self) -> list[tuple[int, int, float]]:
        return list(zip(self.tail, self.head, self.cost))

    def check_invariant(self) -> None:
        n, m = self.n, self.m
        assert self.point[0] == 0 and self.point[n] == m
        assert self.rpoint[0] == 0 and self.rpoint[n] == m
        for i in range(n):
            assert self.point[i] <= self.point[i + 1]
            assert self.rpoint[i] <= self.rpoint[i + 1]
            for e in range(self.point[i], self.point[i + 1]):
                assert self.tail[e] == i
            seg = self.trace[self.rpoint[i] : self.rpoint[i + 1]]
            assert all(seg[k] < seg[k + 1] for k in range(len(seg) - 1))
            assert all(self.head[e] == i for e in seg)
        assert sorted(self.trace) == list(range(m))