Add doctest and fix mypy type annotation in bellman ford

graphs/bellman_ford.py

@@ -1,56 +1,73 @@
 from __future__ import annotations
 
 
-def printDist(dist, V):
-    print("Vertex Distance")
-    distances = ("INF" if d == float("inf") else d for d in dist)
-    print("\t".join(f"{i}\t{d}" for i, d in enumerate(distances)))
+def print_distance(distance: list[float], src):
+    print(f"Vertex\tShortest Distance from vertex {src}")
+    for i, d in enumerate(distance):
+        print(f"{i}\t\t{d}")
 
 
-def BellmanFord(graph: list[dict[str, int]], V: int, E: int, src: int) -> int:
+def check_negative_cycle(
+    graph: list[dict[str, int]], distance: list[float], edge_count: int
+):
+    for j in range(edge_count):
+        u, v, w = [graph[j][k] for k in ["src", "dst", "weight"]]
+        if distance[u] != float("inf") and distance[u] + w < distance[v]:
+            return True
+    return False
+
+
+def bellman_ford(
+    graph: list[dict[str, int]], vertex_count: int, edge_count: int, src: int
+) -> list[float]:
     """
     Returns shortest paths from a vertex src to all
     other vertices.
+    >>> edges = [(2, 1, -10), (3, 2, 3), (0, 3, 5), (0, 1, 4)]
+    >>> g = [{"src": s, "dst": d, "weight": w} for s, d, w in edges]
+    >>> bellman_ford(g, 4, 4, 0)
+    [0.0, -2.0, 8.0, 5.0]
+    >>> g = [{"src": s, "dst": d, "weight": w} for s, d, w in edges + [(1, 3, 5)]]
+    >>> bellman_ford(g, 4, 5, 0)
+    Traceback (most recent call last):
+     ...
+    Exception: Negative cycle found
     """
-    mdist = [float("inf") for i in range(V)]
-    mdist[src] = 0.0
+    distance = [float("inf")] * vertex_count
+    distance[src] = 0.0
 
-    for i in range(V - 1):
-        for j in range(E):
-            u = graph[j]["src"]
-            v = graph[j]["dst"]
-            w = graph[j]["weight"]
+    for i in range(vertex_count - 1):
+        for j in range(edge_count):
+            u, v, w = [graph[j][k] for k in ["src", "dst", "weight"]]
 
-            if mdist[u] != float("inf") and mdist[u] + w < mdist[v]:
-                mdist[v] = mdist[u] + w
-    for j in range(E):
-        u = graph[j]["src"]
-        v = graph[j]["dst"]
-        w = graph[j]["weight"]
+            if distance[u] != float("inf") and distance[u] + w < distance[v]:
+                distance[v] = distance[u] + w
 
-        if mdist[u] != float("inf") and mdist[u] + w < mdist[v]:
-            print("Negative cycle found. Solution not possible.")
-            return
+    negative_cycle_exists = check_negative_cycle(graph, distance, edge_count)
+    if negative_cycle_exists:
+        raise Exception("Negative cycle found")
 
-    printDist(mdist, V)
-    return src
+    return distance
 
 
 if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()
+
     V = int(input("Enter number of vertices: ").strip())
     E = int(input("Enter number of edges: ").strip())
 
-    graph = [dict() for j in range(E)]
+    graph: list[dict[str, int]] = [dict() for j in range(E)]
 
     for i in range(E):
-        graph[i][i] = 0.0
+        print("Edge ", i + 1)
+        src, dest, weight = [
+            int(x)
+            for x in input("Enter source, destination, weight: ").strip().split(" ")
+        ]
+        graph[i] = {"src": src, "dst": dest, "weight": weight}
 
-    for i in range(E):
-        print("\nEdge ", i + 1)
-        src = int(input("Enter source:").strip())
-        dst = int(input("Enter destination:").strip())
-        weight = float(input("Enter weight:").strip())
-        graph[i] = {"src": src, "dst": dst, "weight": weight}
-
-    gsrc = int(input("\nEnter shortest path source:").strip())
-    BellmanFord(graph, V, E, gsrc)
+    source = int(input("\nEnter shortest path source:").strip())
+    shortest_distance = bellman_ford(graph, V, E, source)
+    print_distance(shortest_distance, 0)