Updated check_bipartite_graph_dfs.py

graphs/check_bipartite_graph_dfs.py

@@ -1,34 +1,55 @@
-# Check whether Graph is Bipartite or Not using DFS
+from collections import defaultdict
 
 
-# A Bipartite Graph is a graph whose vertices can be divided into two independent sets,
-# U and V such that every edge (u, v) either connects a vertex from U to V or a vertex
-# from V to U. In other words, for every edge (u, v), either u belongs to U and v to V,
-# or u belongs to V and v to U. We can also say that there is no edge that connects
-# vertices of same set.
-def check_bipartite_dfs(graph):
-    visited = [False] * len(graph)
-    color = [-1] * len(graph)
+def is_bipartite(graph: defaultdict[int, list[int]]) -> bool:
+    """
+    Check whether a graph is Bipartite or not using Depth-First Search (DFS).
 
-    def dfs(v, c):
-        visited[v] = True
-        color[v] = c
-        for u in graph[v]:
-            if not visited[u]:
-                dfs(u, 1 - c)
+    A Bipartite Graph is a graph whose vertices can be divided into two independent
+    sets, U and V such that every edge (u, v) either connects a vertex from
+    U to V or a vertex from V to U. In other words, for every edge (u, v),
+    either u belongs to U and v to V, or u belongs to V and v to U. There is
+    no edge that connects vertices of the same set.
 
-    for i in range(len(graph)):
-        if not visited[i]:
-            dfs(i, 0)
+    Args:
+        graph: An adjacency list representing the graph.
 
-    for i in range(len(graph)):
-        for j in graph[i]:
-            if color[i] == color[j]:
-                return False
+    Returns:
+        True if there's no edge that connects vertices of the same set, False otherwise.
 
-    return True
+    Examples:
+        >>> is_bipartite(
+        ...     defaultdict(list, {0: [1, 2], 1: [0, 3], 2: [0, 4], 3: [1], 4: [2]})
+        ... )
+        False
+        >>> is_bipartite(defaultdict(list, {0: [1, 2], 1: [0, 2], 2: [0, 1]}))
+        True
+    """
 
+    def depth_first_search(node: int, color: int) -> bool:
+        visited[node] = color
+        return any(
+            visited[neighbour] == color
+            or (
+                visited[neighbour] == -1
+                and not depth_first_search(neighbour, 1 - color)
+            )
+            for neighbour in graph[node]
+        )
 
-# Adjacency list of graph
-graph = {0: [1, 3], 1: [0, 2], 2: [1, 3], 3: [0, 2], 4: []}
-print(check_bipartite_dfs(graph))
+    visited: defaultdict[int, int] = defaultdict(lambda: -1)
+
+    return all(
+        not (visited[node] == -1 and not depth_first_search(node, 0)) for node in graph
+    )
+
+
+if __name__ == "__main__":
+    import doctest
+
+    result = doctest.testmod()
+
+    if result.failed:
+        print(f"{result.failed} test(s) failed.")
+    else:
+        print("All tests passed!")