update graphs/breadth_first_search.py

graphs/breadth_first_search.py

@@ -2,33 +2,35 @@
 
 """ Author: OMKAR PATHAK """
 
+from typing import Set
+
 
 class Graph:
     def __init__(self):
         self.vertices = {}
 
-    def printGraph(self):
+    def print_graph(self) -> None:
         """prints adjacency list representation of graaph"""
-        for i in self.vertices.keys():
+        for i in self.vertices:
             print(i, " : ", " -> ".join([str(j) for j in self.vertices[i]]))
 
-    def addEdge(self, fromVertex, toVertex):
+    def add_edge(self, from_vertex: int, to_vertex: int) -> None:
         """adding the edge between two vertices"""
-        if fromVertex in self.vertices.keys():
-            self.vertices[fromVertex].append(toVertex)
+        if from_vertex in self.vertices:
+            self.vertices[from_vertex].append(to_vertex)
         else:
-            self.vertices[fromVertex] = [toVertex]
+            self.vertices[from_vertex] = [to_vertex]
 
-    def BFS(self, startVertex):
+    def bfs(self, start_vertex: int) -> Set[int]:
         # initialize set for storing already visited vertices
         visited = set()
 
         # create a first in first out queue to store all the vertices for BFS
         queue = []
 
         # mark the source node as visited and enqueue it
-        visited.add(startVertex)
-        queue.append(startVertex)
+        visited.add(start_vertex)
+        queue.append(start_vertex)
 
         while queue:
             vertex = queue.pop(0)
@@ -43,17 +45,17 @@ def BFS(self, startVertex):
 
 if __name__ == "__main__":
     g = Graph()
-    g.addEdge(0, 1)
-    g.addEdge(0, 2)
-    g.addEdge(1, 2)
-    g.addEdge(2, 0)
-    g.addEdge(2, 3)
-    g.addEdge(3, 3)
-
-    g.printGraph()
+    g.add_edge(0, 1)
+    g.add_edge(0, 2)
+    g.add_edge(1, 2)
+    g.add_edge(2, 0)
+    g.add_edge(2, 3)
+    g.add_edge(3, 3)
+
+    g.print_graph()
     # 0  :  1 -> 2
     # 1  :  2
     # 2  :  0 -> 3
     # 3  :  3
 
-    assert sorted(g.BFS(2)) == [0, 1, 2, 3]
+    assert sorted(g.bfs(2)) == [0, 1, 2, 3]