Add doctest and fix mypy type annotation in bellman ford

Author: mhihasan

graphs/bellman_ford.py

@@ -1,56 +1,53 @@
-from __future__ import annotations
+from typing import Dict, List
 
 
-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 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]["src"], graph[j]["dst"], graph[j]["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
+    for j in range(edge_count):
+        u, v, w = graph[j]["src"], graph[j]["dst"], graph[j]["weight"]
 
-    printDist(mdist, V)
-    return src
+        if distance[u] != float("inf") and distance[u] + w < distance[v]:
+            raise Exception("Negative cycle found")
 
+    return distance
 
-if __name__ == "__main__":
-    V = int(input("Enter number of vertices: ").strip())
-    E = int(input("Enter number of edges: ").strip())
-
-    graph = [dict() for j in range(E)]
 
-    for i in range(E):
-        graph[i][i] = 0.0
+if __name__ == "__main__":
+    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]
+    shortest_distance = bellman_ford(g, 4, 4, 0)
+    print_distance(shortest_distance, 0)
 
-    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}
+    import doctest
 
-    gsrc = int(input("\nEnter shortest path source:").strip())
-    BellmanFord(graph, V, E, gsrc)
+    doctest.testmod()