Eppstein's k shortest paths in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/eppstein-k-shortest-paths/. The authors are listed in prov/WNETNH.jsonl.

import heapq
import math

INF = math.inf


class _Node:
    """A node of a persistent leftist heap. It holds the best sidetrack out of one vertex."""

    __slots__ = ("key", "vertex", "left", "right", "rank")

    def __init__(self, key, vertex, left=None, right=None):
        self.key = key
        self.vertex = vertex
        self.left = left
        self.right = right
        self.rank = 1 + min(_rank(left), _rank(right))


def _rank(x) -> int:
    return x.rank if x else 0


def _insert(h, node):
    """Persistent insertion of a fresh one-node heap: h is left unchanged.

    Only the nodes on the right spine of h are copied, O(log size) of them.
    """
    if h is None:
        return node
    if node.key < h.key:
        return _Node(node.key, node.vertex, h, None)
    left, right = h.left, _insert(h.right, node)
    if _rank(left) < _rank(right):
        left, right = right, left
    return _Node(h.key, h.vertex, left, right)


def _relevant(adj, radj, s, t):
    """Vertices on some s-t walk: reachable from s and reaching t."""
    n = len(adj)

    def reach(g, r):
        seen = [False] * n
        seen[r] = True
        stack = [r]
        while stack:
            u = stack.pop()
            for v, *_ in g[u]:
                if not seen[v]:
                    seen[v] = True
                    stack.append(v)
        return seen

    a, b = reach(adj, s), reach(radj, t)
    return [a[v] and b[v] for v in range(n)]


def _distances_to_t(adj, radj, t, keep):
    """d[v] = dist(v, t) inside the kept vertices, and the tree edge out of each v.

    The tree edge of v is the edge (v, nxt[v]) with index tree_edge[v] in adj[v] that
    last lowered d[v]. Returns None if a negative cycle lies among the kept vertices.
    """
    n = len(adj)
    d = [INF] * n
    nxt = [-1] * n
    tree_edge = [-1] * n
    d[t] = 0
    if all(w >= 0 for u in range(n) if keep[u] for _, w in adj[u]):
        heap = [(0, t)]
        done = [False] * n
        while heap:
            dv, v = heapq.heappop(heap)
            if done[v]:
                continue
            done[v] = True
            for u, w, i in radj[v]:
                if keep[u] and dv + w < d[u]:
                    d[u], nxt[u], tree_edge[u] = dv + w, v, i
                    heapq.heappush(heap, (d[u], u))
        return d, nxt, tree_edge
    for _ in range(n):  # Bellman-Ford on the reversed graph
        changed = False
        for v in range(n):
            if not keep[v] or d[v] == INF:
                continue
            for u, w, i in radj[v]:
                if keep[u] and d[v] + w < d[u]:
                    d[u], nxt[u], tree_edge[u] = d[v] + w, v, i
                    changed = True
        if not changed:
            return d, nxt, tree_edge
    return None


def eppstein(
    adj: list[list[tuple[int, float]]], s: int, t: int, k: int
) -> list[tuple[float, list[int]]]:
    """Return the k shortest s-t walks as (weight, vertices), lightest first.

    The graph has the vertices 0 .. n-1; adj[u] lists the edges (v, w) out of u.
    Weights are real numbers, negative ones allowed. Walks may repeat vertices and
    edges. If a negative cycle lies on some s-t walk, there is no shortest walk and
    ValueError is raised; negative cycles elsewhere are ignored. If t is not
    reachable from s, the list is empty; if fewer than k walks exist, all are returned.
    """
    n = len(adj)
    radj: list[list[tuple[int, float, int]]] = [[] for _ in range(n)]
    for u in range(n):
        for i, (v, w) in enumerate(adj[u]):
            radj[v].append((u, w, i))
    keep = _relevant(adj, radj, s, t)
    if not keep[s] or k <= 0:
        return []
    found = _distances_to_t(adj, radj, t, keep)  # @line L1
    if found is None:
        raise ValueError("a negative cycle lies on an s-t walk")
    d, nxt, tree_edge = found
    # sidetracks out of each vertex, as a binary heap array; the best one is first
    side: list[list[tuple[float, int]]] = [[] for _ in range(n)]
    for u in range(n):
        if keep[u]:
            for i, (v, w) in enumerate(adj[u]):
                if keep[v] and i != tree_edge[u]:
                    side[u].append((w + d[v] - d[u], i))  # @line L2
            heapq.heapify(side[u])
    # H[v]: persistent heap of the best sidetracks of v and of its tree ancestors
    hg: list[_Node | None] = [None] * n
    kids: list[list[int]] = [[] for _ in range(n)]
    for v in range(n):
        if nxt[v] != -1:
            kids[nxt[v]].append(v)
    order = [t]  # tree order from t: the tree parent comes first
    for v in order:
        order.extend(kids[v])
    for v in order:
        h = hg[nxt[v]] if v != t else None
        if side[v]:
            h = _insert(h, _Node(side[v][0][0], v))  # @line L3
        hg[v] = h
    result = [(d[s], _walk(adj, nxt, s, t, []))]  # @line L4
    pq: list = []
    tick = 0
    if hg[s] is not None:
        pq.append((d[s] + hg[s].key, tick, ("T", hg[s]), ()))  # @line L5
    while pq and len(result) < k:  # @line L6
        cost, _, item, hist = heapq.heappop(pq)  # @line L7
        key, u, i, children = _expand(item, side)
        seq = hist + ((u, i),)
        result.append((cost, _walk(adj, nxt, s, t, list(seq))))  # @line L8
        for child in children:  # @line L9
            ckey = _expand(child, side)[0]
            tick += 1
            heapq.heappush(pq, (cost - key + ckey, tick, child, hist))
        head = hg[adj[u][i][0]]
        if head is not None:  # @line L10
            tick += 1
            heapq.heappush(pq, (cost + head.key, tick, ("T", head), seq))
    return result  # @line L11


def _expand(item, side):
    """Key, sidetrack (u, edge index) and path-graph heap children of a node."""
    if item[0] == "T":
        node = item[1]
        key, i = side[node.vertex][0]
        children = [("T", c) for c in (node.left, node.right) if c is not None]
        # the rest of the sidetracks of the vertex hang below the T node
        children += [("A", node.vertex, c) for c in (1, 2) if c < len(side[node.vertex])]
        return key, node.vertex, i, children
    _, u, j = item  # position j >= 1 of the heap array of the sidetracks of u
    key, i = side[u][j]
    kids = [("A", u, c) for c in (2 * j + 1, 2 * j + 2) if c < len(side[u])]
    return key, u, i, kids


def _walk(adj, nxt, s: int, t: int, sidetracks: list) -> list[int]:
    """Follow tree edges from s, taking the given sidetracks in order, to t."""
    walk = [s]
    v = s
    for u, i in sidetracks:
        while v != u:
            v = nxt[v]
            walk.append(v)
        v = adj[u][i][0]
        walk.append(v)
    while v != t:
        v = nxt[v]
        walk.append(v)
    return walk