Updated check_bipartite_graph_dfs.py

graphs/check_bipartite_graph_dfs.py

@@ -1,34 +1,50 @@
-# Check whether Graph is Bipartite or Not using DFS
-
-
-# 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 dfs(v, c):
-        visited[v] = True
-        color[v] = c
-        for u in graph[v]:
-            if not visited[u]:
-                dfs(u, 1 - c)
-
-    for i in range(len(graph)):
-        if not visited[i]:
-            dfs(i, 0)
-
-    for i in range(len(graph)):
-        for j in graph[i]:
-            if color[i] == color[j]:
+from collections import defaultdict
+
+
+def is_bipartite(graph: defaultdict[int, list[int]]) -> bool:
+    """
+    Check whether a graph is Bipartite or not using Depth-First Search (DFS).
+
+    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.
+
+    Args:
+        graph: An adjacency list representing the graph.
+
+    Returns:
+        True if there's no edge that connects vertices of the same set, False otherwise.
+
+    Examples:
+        >>> graph1=defaultdict(list,{0: [1, 3], 1: [0, 2], 2: [1, 3], 3: [0, 2], 4: []})
+        >>> is_bipartite(graph1)
+        True
+
+        >>> graph2 = defaultdict(list, {0: [1, 2], 1: [0, 2], 2: [0, 1]})
+        >>> is_bipartite(graph2)
+        False
+    """
+
+    def dfs(node, color):
+        vis[node] = color
+        for nbor in graph[node]:
+            if vis[nbor] == color or (vis[nbor] == -1 and not dfs(nbor, 1 - color)):
                 return False
+        return True
+
+    vis: defaultdict[int, int] = defaultdict(lambda: -1)
+
+    return all(not (vis[node] == -1 and not dfs(node, 0)) for node in graph)
+
 
-    return True
+if __name__ == "__main__":
+    import doctest
 
+    result = doctest.testmod()
 
-# Adjacency list of graph
-graph = {0: [1, 3], 1: [0, 2], 2: [1, 3], 3: [0, 2], 4: []}
-print(check_bipartite_dfs(graph))
+    if result.failed == 0:
+        print("All tests passed!")
+    else:
+        print(f"{result.failed} test(s) failed.")