Added tests

Author: joelkurien

graphs/johnson_graph.py

@@ -9,30 +9,58 @@ class JohnsonGraph:
     def __init__(self) -> None:
         """
         Initializes an empty graph with no edges.
+        >>> g = JohnsonGraph()
+        >>> g.edges
+        []
+        >>> g.graph
+        {}
         """
         self.edges: list[tuple[str, str, int]] = []
         self.graph: dict[str, list[tuple[str, int]]] = {}
 
     # add vertices for a graph
     def add_vertices(self, vertex: str) -> None:
         """
-        Adds a vertex `u` to the graph with an empty adjacency list.
+        Adds a vertex `vertex` to the graph with an empty adjacency list.
+        >>> g = JohnsonGraph()
+        >>> g.add_vertices("A")
+        >>> g.graph
+        {'A': []}
         """
         self.graph[vertex] = []
 
     # assign weights for each edges formed of the directed graph
     def add_edge(self, vertex_a: str, vertex_b: str, weight: int) -> None:
         """
-        Adds a directed edge from vertex `u` to vertex `v` with weight `w`.
+        Adds a directed edge from vertex `vertex_a` 
+        to vertex `vertex_b` with weight `weight`.
+        >>> g = JohnsonGraph()
+        >>> g.add_vertices("A")
+        >>> g.add_vertices("B")
+        >>> g.add_edge("A", "B", 5)
+        >>> g.edges
+        [('A', 'B', 5)]
+        >>> g.graph
+        {'A': [('B', 5)], 'B': []}
         """
         self.edges.append((vertex_a, vertex_b, weight))
         self.graph[vertex_a].append((vertex_b, weight))
 
     # perform a dijkstra algorithm on a directed graph
     def dijkstra(self, start: str) -> dict:
         """
-        Computes the shortest path from vertex `s`
+        Computes the shortest path from vertex `start` 
         to all other vertices using Dijkstra's algorithm.
+        >>> g = JohnsonGraph()
+        >>> g.add_vertices("A")
+        >>> g.add_vertices("B")
+        >>> g.add_edge("A", "B", 1)
+        >>> g.dijkstra("A")
+        {'A': 0, 'B': 1}
+        >>> g.add_vertices("C")
+        >>> g.add_edge("B", "C", 2)
+        >>> g.dijkstra("A")
+        {'A': 0, 'B': 1, 'C': 3}
         """
         distances = {vertex: sys.maxsize - 1 for vertex in self.graph}
         pq = [(0, start)]
@@ -52,8 +80,18 @@ def dijkstra(self, start: str) -> dict:
     # carry out the bellman ford algorithm for a node and estimate its distance vector
     def bellman_ford(self, start: str) -> dict:
         """
-        Computes the shortest path from vertex `s`
-        to all other vertices using the Bellman-Ford algorithm.
+        Computes the shortest path from vertex `start` to 
+        all other vertices using the Bellman-Ford algorithm.
+        >>> g = JohnsonGraph()
+        >>> g.add_vertices("A")
+        >>> g.add_vertices("B")
+        >>> g.add_edge("A", "B", 1)
+        >>> g.bellman_ford("A")
+        {'A': 0, 'B': 1}
+        >>> g.add_vertices("C")
+        >>> g.add_edge("B", "C", 2)
+        >>> g.bellman_ford("A")
+        {'A': 0, 'B': 1, 'C': 3}
         """
         distances = {vertex: sys.maxsize - 1 for vertex in self.graph}
         distances[start] = 0
@@ -73,8 +111,19 @@ def bellman_ford(self, start: str) -> dict:
     # or the bellman ford algorithm efficiently
     def johnson_algo(self) -> list[dict]:
         """
-        Computes the shortest paths between
-        all pairs of vertices using Johnson's algorithm.
+        Computes the shortest paths between 
+        all pairs of vertices using Johnson's algorithm
+        for a directed graph.
+        >>> g = JohnsonGraph()
+        >>> g.add_vertices("A")
+        >>> g.add_vertices("B")
+        >>> g.add_vertices("C")
+        >>> g.add_edge("A", "B", 1)
+        >>> g.add_edge("B", "C", 2)
+        >>> g.add_edge("A", "C", 4)
+        >>> optimal_paths = g.johnson_algo()
+        >>> optimal_paths
+        [{'A': 0, 'B': 1, 'C': 3}, {'A': None, 'B': 0, 'C': 2}, {'A': None, 'B': None, 'C': 0}]
         """
         self.add_vertices("#")
         for vertex in self.graph:
@@ -95,36 +144,26 @@ def johnson_algo(self) -> list[dict]:
                             weight + hash_path[vertex_a] - hash_path[vertex_b])
 
         self.graph.pop("#")
-        self.edges = [
-            (vertex1, vertex2, node_weight)
-            for vertex1, vertex2, node_weight in self.edges
-            if vertex1 != "#"
-        ]
         filtered_edges = []
         for vertex1, vertex2, node_weight in self.edges:
-            if vertex1 != "#":
-                filtered_edges.append((vertex1, vertex2, node_weight))
+            filtered_edges.append((vertex1, vertex2, node_weight))
         self.edges = filtered_edges
 
         for vertex in self.graph:
-            self.graph[vertex] = [
-                (vertex2, node_weight)
-                for vertex1, vertex2, node_weight in self.edges
-                if vertex1 == vertex
-            ]
-
-            filtered_neighbors = []
+            self.graph[vertex] = []
             for vertex1, vertex2, node_weight in self.edges:
                 if vertex1 == vertex:
-                    filtered_neighbors.append((vertex2, node_weight))
-            self.graph[vertex] = filtered_neighbors
+                    self.graph[vertex].append((vertex2, node_weight))
 
         distances = []
         for vertex1 in self.graph:
             new_dist = self.dijkstra(vertex1)
             for vertex2 in self.graph:
-                if new_dist[vertex2] < sys.maxsize - 1:
-                    new_dist[vertex2] += hash_path[vertex1] - hash_path[vertex2]
+                if new_dist[vertex2] < sys.maxsize-1:
+                    new_dist[vertex2] += hash_path[vertex2] - hash_path[vertex1]
+            for key in new_dist:
+                if new_dist[key] == sys.maxsize-1:
+                    new_dist[key] = None
             distances.append(new_dist)
         return distances