Bellman-Ford Shortest Paths With Negative Cycle

from __future__ import annotations

from dataclasses import dataclass


@dataclass(frozen=True)
class Edge:
    source: str
    target: str
    weight: float


def bellman_ford(
    nodes: list[str],
    edges: list[Edge],
    start: str,
) -> tuple[dict[str, float], dict[str, str | None], list[str] | None]:
    distance = {node: float("inf") for node in nodes}
    previous: dict[str, str | None] = {node: None for node in nodes}
    distance[start] = 0.0

    for _ in range(len(nodes) - 1):
        relaxed = False
        for edge in edges:
            candidate = distance[edge.source] + edge.weight
            if candidate < distance[edge.target]:
                distance[edge.target] = candidate
                previous[edge.target] = edge.source
                relaxed = True
        if not relaxed:
            break

    return distance, previous, find_negative_cycle(nodes, edges, distance, previous)


def find_negative_cycle(
    nodes: list[str],
    edges: list[Edge],
    distance: dict[str, float],
    previous: dict[str, str | None],
) -> list[str] | None:
    for edge in edges:
        if distance[edge.source] + edge.weight < distance[edge.target]:
            previous[edge.target] = edge.source
            node = edge.target
            for _ in range(len(nodes)):
                node = previous[node]  # type: ignore[assignment]
            cycle = [node]
            cursor = previous[node]
            while cursor != node and cursor is not None:
                cycle.append(cursor)
                cursor = previous[cursor]
            cycle.append(node)
            cycle.reverse()
            return cycle
    return None


def shortest_path(previous: dict[str, str | None], start: str, target: str) -> list[str]:
    path = [target]
    cursor = target
    while cursor != start:
        parent = previous[cursor]
        if parent is None:
            return []
        path.append(parent)
        cursor = parent
    path.reverse()
    return path


def main() -> None:
    nodes = ["A", "B", "C", "D"]
    edges = [
        Edge("A", "B", 4.0),
        Edge("A", "C", 5.0),
        Edge("B", "C", -2.5),
        Edge("C", "D", 3.0),
        Edge("B", "D", 6.0),
    ]
    distance, previous, cycle = bellman_ford(nodes, edges, "A")
    print("cycle:", cycle)
    print("D cost:", distance["D"])
    print("path:", " -> ".join(shortest_path(previous, "A", "D")))

    risky = [Edge("A", "B", 1.0), Edge("B", "C", -1.0), Edge("C", "B", -1.0)]
    _, _, cycle = bellman_ford(["A", "B", "C"], risky, "A")
    print("negative cycle:", cycle)


if __name__ == "__main__":
    main()