Contraction hierarchies in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/contraction-hierarchies/. The authors are listed in prov/K2PFVE.jsonl.

import heapq
import math

INF = math.inf


class CH:
    """A contraction hierarchy: ranks, upward and downward edges, shortcut middles."""

    def __init__(self, n: int) -> None:
        self.rank = [0] * n
        self.up: list[list[tuple[int, float]]] = [[] for _ in range(n)]
        self.down: list[list[tuple[int, float]]] = [[] for _ in range(n)]
        self.mid: dict[tuple[int, int], int | None] = {}
        self.shortcuts = 0


def _witness(out, u: int, x: int, skip: int, limit: float) -> bool:
    """True if the remaining graph without skip has a u-x path of weight <= limit."""
    d = {u: 0}
    heap = [(0, u)]
    while heap:
        du, a = heapq.heappop(heap)
        if du > limit:
            return False
        if du > d[a]:
            continue
        if a == x:
            return True
        for b, (w, _) in out[a].items():
            if b != skip and du + w < d.get(b, INF):
                d[b] = du + w
                heapq.heappush(heap, (d[b], b))
    return False


def _needed(out, inn, v: int) -> list[tuple[int, int, float]]:
    """The shortcuts that contracting v would add: (u, x, weight)."""
    res = []
    for u, (wu, _) in inn[v].items():
        for x, (wx, _) in out[v].items():  # @line L2
            if u != x and not _witness(out, u, x, v, wu + wx):  # @line L3
                res.append((u, x, wu + wx))
    return res


def ch_preprocess(
    adj: list[list[tuple[int, float]]], order: list[int] | None = None
) -> CH:
    """Build a contraction hierarchy of the directed graph adj.

    The graph has the vertices 0 .. n-1; adj[u] lists the edges (v, w) out of u.
    Precondition: every weight is nonnegative (a ValueError is raised otherwise).
    If order is given, vertices are contracted in that order; otherwise by the
    edge-difference heuristic with lazy updates. Witness searches are exact.
    """
    n = len(adj)
    out: list[dict[int, tuple[float, int | None]]] = [{} for _ in range(n)]
    inn: list[dict[int, tuple[float, int | None]]] = [{} for _ in range(n)]
    for u in range(n):
        for v, w in adj[u]:
            if w < 0:
                raise ValueError("contraction hierarchies need nonnegative weights")
            if u != v and w < out[u].get(v, (INF, None))[0]:
                out[u][v] = (w, None)
                inn[v][u] = (w, None)
    ch = CH(n)

    def priority(v: int) -> int:
        return len(_needed(out, inn, v)) - len(out[v]) - len(inn[v])

    if order is None:
        heap = [(priority(v), v) for v in range(n)]
        heapq.heapify(heap)
        seq = []
        done = [False] * n
        while heap:
            p, v = heapq.heappop(heap)
            if done[v]:
                continue
            q = priority(v)
            if heap and q > heap[0][0]:
                heapq.heappush(heap, (q, v))
                continue
            done[v] = True
            seq.append(v)
            _contract(ch, out, inn, v, len(seq) - 1)
        return ch
    for i, v in enumerate(order):  # @line L1
        _contract(ch, out, inn, v, i)
    return ch


def _contract(ch: CH, out, inn, v: int, r: int) -> None:
    ch.rank[v] = r
    for u, x, w in _needed(out, inn, v):
        if w < out[u].get(x, (INF, None))[0]:
            out[u][x] = (w, v)  # @line L4
            inn[x][u] = (w, v)
            ch.shortcuts += 1
    for x, (w, m) in out[v].items():  # @line L5
        ch.up[v].append((x, w))
        ch.mid[(v, x)] = m
        del inn[x][v]
    for u, (w, m) in inn[v].items():
        ch.down[v].append((u, w))
        ch.mid[(u, v)] = m
        del out[u][v]
    out[v].clear()
    inn[v].clear()


def _unpack(ch: CH, u: int, v: int, acc: list[int]) -> None:
    m = ch.mid[(u, v)]
    if m is None:
        acc.append(v)
    else:
        _unpack(ch, u, m, acc)
        _unpack(ch, m, v, acc)


def ch_query(ch: CH, s: int, t: int) -> tuple[float, list[int] | None]:
    """Return (dist, path) from s to t, or (math.inf, None) if t is unreachable."""
    if s == t:
        return 0, [s]
    n = len(ch.rank)
    dist = [[INF] * n, [INF] * n]
    pred: list[list[int | None]] = [[None] * n, [None] * n]
    done = [[False] * n, [False] * n]
    graph = [ch.up, ch.down]
    dist[0][s] = 0  # @line L6
    dist[1][t] = 0
    heaps = [[(0, s)], [(0, t)]]
    mu, meet = INF, None
    while True:
        side = min((0, 1), key=lambda i: heaps[i][0][0] if heaps[i] else INF)
        if not heaps[side] or heaps[side][0][0] >= mu:  # @line L7
            break
        du, u = heapq.heappop(heaps[side])
        if done[side][u]:
            continue
        done[side][u] = True
        if du + dist[1 - side][u] < mu:  # @line L8
            mu, meet = du + dist[1 - side][u], u
        for v, w in graph[side][u]:
            if du + w < dist[side][v]:
                dist[side][v] = du + w
                pred[side][v] = u
                heapq.heappush(heaps[side], (dist[side][v], v))
    if meet is None:  # @line L9
        return INF, None
    fwd = [meet]
    while fwd[-1] != s:
        fwd.append(pred[0][fwd[-1]])
    fwd.reverse()
    bwd = [meet]
    while bwd[-1] != t:
        bwd.append(pred[1][bwd[-1]])
    up_path = fwd + bwd[1:]
    path = [s]
    for a, b in zip(up_path, up_path[1:]):
        _unpack(ch, a, b, path)  # @line L10
    return mu, path