Update graph.py

Author: ivanmmarkovic

graphs/kosarajus-algorithm/graph.py

@@ -6,63 +6,66 @@
 
 class Graph:
 
+
     def __init__(self) -> None:
-        self.vertices: List[str] = []
-        self.colors: Dict[str, str] = {}
-        self.adjacency_list: Dict[str, Set[str]] = {}
-        self.adjacency_list_reversed: Dict[str, Set[str]] = {}
-        self.stack: Stack = Stack()
-        self.components: List[List[str]] = []
-
-    def add_vertex(self, label:str):
-        self.vertices.append(label)
-        self.colors[label] = 'white'
+        self.vertices:Set[str] = set()
+        self.adjacency_list:Dict[str, Set[str]] = dict()
+        self.adjacency_list_reversed:Dict[str, Set[str]] = dict()
+        self.visited:Set[str] = set()
+        self.stack:Stack = Stack()
+
+
+    def add_vertex(self, label:str) -> None:
+        self.vertices.add(label)
         self.adjacency_list[label] = set()
         self.adjacency_list_reversed[label] = set()
 
-    def add_edge(self, label1:str, label2:str):
+
+    def add_edge(self, label1:str, label2:str) -> None:
+        if label1 not in self.vertices or label2 not in self.vertices:
+            raise Exception('Vertices are not added')
         self.adjacency_list[label1].add(label2)
         self.adjacency_list_reversed[label2].add(label1)
 
-    def kosaraju(self):
-        for label in self.vertices:
-            if self.colors[label] == 'white':
-                self.dfs(label)
 
-        for color in self.colors:
-            self.colors[color] = 'white'
+    def kosaraju(self) -> List[List[str]]:
+        for v in self.vertices:
+            if v not in self.visited:
+                self._dfs(v)
+
+        self.visited.clear()
+
+        connected_components:List[List[str]] = list()
 
         while not self.stack.is_empty():
-            current:str = self.stack.pop()
+            v:str = self.stack.pop()
+            if v not in self.visited:
+                connected:List[str] = list()
+                self._dfs_reversed(v, connected)
 
-            if self.colors[current] == 'white':
-                connected_components: List[str] = []
-                self.dfs_reversed(current, connected_components)
+                if len(connected) > 0:
+                    connected_components.append(connected)
 
-                self.components.append(connected_components)
+        return list(connected_components)
 
-        print(self.components)
 
 
-    def dfs(self, label:str):
-        self.colors[label] = 'gray'
+    def _dfs(self, label:str) -> None:
+        self.visited.add(label)
 
         for n in self.adjacency_list[label]:
-            if self.colors[n] == 'white':
-                self.dfs(n)
+            if n not in self.visited:
+                self._dfs(n)
 
-        self.colors[label] = 'black'
         self.stack.push(label)
 
-    def dfs_reversed(self, label: str, connected_components:List[str]):
-        self.colors[label] = 'gray'
-
-        for n in self.adjacency_list_reversed[label]:
-            if self.colors[n] == 'white':
-                self.dfs_reversed(n, connected_components)
-        connected_components.append(label)
-        self.colors[label] = 'black'
 
+    def _dfs_reversed(self, v:str, connected:List[str]) -> None:
+        connected.append(v)
+        self.visited.add(v)
+        for n in self.adjacency_list_reversed[v]:
+            if n not in self.visited:
+                self._dfs_reversed(n, connected)