Improve dijkstra prints

Author: andrelbd1

solutions/dijkstra/dijkstra.py

@@ -1,24 +1,26 @@
 from math import inf
 
+
 class Dijkstra:
     def __init__(self, graph):
         self.graph = graph
         self.n_vertices = len(graph)
 
-    def print_distance(self, dist, target=None):
+    def print_distance(self, dist, target=None, show_print=None):
         if target:
-            print("distance:", dist[target])
+            if show_print:
+                print("distance:", dist[target])
             return dist[target]
-
-        print("Vertex \t Distance from Source")
-        for node in range(self.n_vertices):
-            print(node, "\t\t", dist[node])
+        if show_print:
+            print("Vertex \t Distance from Source")
+            for node in range(self.n_vertices):
+                print(node, "\t\t", dist[node])
 
     def print_path(self, path, target):
         queue = [target]
         build_path = []
         while True:
-            prev=queue.pop(0)
+            prev = queue.pop(0)
             build_path.append(prev)
             if path[prev] is None:
                 break
@@ -27,33 +29,30 @@ def print_path(self, path, target):
         print("path:", " -> ".join([str(i) for i in build_path]))
 
     def min_distance(self, dist, visited):
-        min = inf # Initialize minimum distance for next node 
-        for v in range(self.n_vertices): # get nearest vertex not visited
-            if dist[v] < min and visited[v] == False:
+        min = inf  # Initialize minimum distance for next node
+        for v in range(self.n_vertices):  # get nearest vertex not visited
+            if dist[v] < min and visited[v] is False:
                 min = dist[v]
                 min_index = v
-
         return min_index
 
-    def dijkstra(self, src, target=None): # O(n2)
+    def dijkstra(self, src, target=None, show_print=None):  # O(n2)
         visited = [False] * self.n_vertices
         dist = [inf] * self.n_vertices
         dist[src] = 0
-        path_prev = {src:None}
-
+        path_prev = {src: None}
         for _ in range(self.n_vertices):
             # Get the minimum distance vertex not visited yet.
             # u is always equal to src in first iteration.
             u = self.min_distance(dist, visited)
             visited[u] = True
-
             # Update distance value of the adjacent vertices
             for v in range(self.n_vertices):
-                if (self.graph[u][v] > 0 and # Check if has edge
-                        visited[v] == False and
+                if (self.graph[u][v] > 0 and  # Check if has edge
+                        visited[v] is False and
                         dist[v] > dist[u] + self.graph[u][v]):
                     dist[v] = dist[u] + self.graph[u][v]
                     path_prev.update({v: u})
-        
-        self.print_path(path_prev, target)
-        return self.print_distance(dist, target)
+        if show_print:
+            self.print_path(path_prev, target)
+        return self.print_distance(dist, target, show_print)

solutions/dijkstra/instance.py

@@ -11,6 +11,5 @@
          [0, 0, 2, 0, 0, 0, 6, 7, 0]]
 
 instance = Dijkstra(graph)
-
-assert instance.dijkstra(src=0,target=4) == 21
-assert instance.dijkstra(src=1,target=5) == 12
+assert instance.dijkstra(src=0, target=4, show_print=True) == 21
+assert instance.dijkstra(src=1, target=5, show_print=True) == 12