Floyd-Warshall All-Pairs Shortest Paths

INF = float("inf")


def floyd_warshall(
    graph: dict[str, dict[str, float]]
) -> tuple[dict[str, dict[str, float]], dict[str, dict[str, str | None]]]:
    nodes = sorted(graph)
    dist = {source: {target: INF for target in nodes} for source in nodes}
    nxt: dict[str, dict[str, str | None]] = {
        source: {target: None for target in nodes} for source in nodes
    }
    for node in nodes:
        dist[node][node] = 0.0
    for source, edges in graph.items():
        for target, weight in edges.items():
            dist[source][target] = weight
            nxt[source][target] = target

    for middle in nodes:
        for source in nodes:
            if dist[source][middle] == INF:
                continue
            for target in nodes:
                candidate = dist[source][middle] + dist[middle][target]
                if candidate < dist[source][target]:
                    dist[source][target] = candidate
                    nxt[source][target] = nxt[source][middle]
    return dist, nxt


def build_path(
    nxt: dict[str, dict[str, str | None]], source: str, target: str
) -> list[str] | None:
    if source == target:
        return [source]
    if nxt[source][target] is None:
        return None
    route = [source]
    while source != target:
        step = nxt[source][target]
        if step is None:
            return None
        source = step
        route.append(source)
    return route


if __name__ == "__main__":
    graph = {
        "A": {"B": 3.0, "D": 7.0},
        "B": {"A": 8.0, "C": 2.0},
        "C": {"A": 5.0, "D": 1.0},
        "D": {"A": 2.0},
    }
    dist, nxt = floyd_warshall(graph)
    print("A to C:", dist["A"]["C"], build_path(nxt, "A", "C"))
    print("D to C:", dist["D"]["C"], build_path(nxt, "D", "C"))
    print("C to C:", dist["C"]["C"], build_path(nxt, "C", "C"))