Correct implementation and add tests for dfs and bfs

Author: prak74

graphs/directed_and_undirected_weighted_graph.py

@@ -14,6 +14,18 @@ def __init__(self):
     # adding the weight is optional
     # handles repetition
     def add_pair(self, u, v, w=1):
+        """
+        Adds a directed edge u->v with weight w.
+        >>> dg = DirectedGraph()
+        >>> dg.add_pair(-1,2)
+        >>> dg.add_pair(1,3,5)
+        >>> dg.add_pair(1,3,5)
+        >>> dg.add_pair(1,3,6)
+        >>> dg.all_nodes()
+        [-1, 2, 1, 3]
+        >>> dg.graph[1]
+        [[5, 3], [6, 3]]
+        """
         if self.graph.get(u):
             if self.graph[u].count([w, v]) == 0:
                 self.graph[u].append([w, v])
@@ -23,53 +35,94 @@ def add_pair(self, u, v, w=1):
             self.graph[v] = []
 
     def all_nodes(self):
+        """
+        Returns list of all nodes in the graph.
+        >>> dg = DirectedGraph()
+        >>> dg.all_nodes()
+        []
+        >>> dg.add_pair(1,1)
+        >>> dg.all_nodes()
+        [1]
+        >>> dg.add_pair(2,3,3)
+        >>> dg.all_nodes()
+        [1, 2, 3]
+        """
         return list(self.graph)
 
     # handles if the input does not exist
     def remove_pair(self, u, v):
+        """
+        Removes all edges u->v if it exists.
+        >>> dg = DirectedGraph()
+        >>> dg.remove_pair(1,2) # silently exits
+        >>> dg.add_pair(0,5,2)
+        >>> dg.graph[0]
+        [[2, 5]]
+        >>> dg.remove_pair(5,0)
+        >>> dg.graph[0]
+        [[2, 5]]
+        >>> dg.remove_pair(0,5)
+        >>> dg.graph[0]
+        []
+        """
         if self.graph.get(u):
             for _ in self.graph[u]:
                 if _[1] == v:
                     self.graph[u].remove(_)
 
     # if no destination is meant the default value is -1
     def dfs(self, s=-2, d=-1):
-        if s == d:
-            return []
+        """
+        Performs depth first search from s to find d.
+        Returns the path s->d as a list.
+        Returns dfs from s if d is not found
+        >>> dg = DirectedGraph()
+        >>> dg.dfs()
+        []
+        >>> dg.add_pair(1,1)
+        >>> dg.dfs(1,1)
+        [1]
+        >>> dg = DirectedGraph()
+        >>> dg.add_pair(0,1)
+        >>> dg.add_pair(0,2)
+        >>> dg.add_pair(1,3)
+        >>> dg.add_pair(1,4)
+        >>> dg.add_pair(1,5)
+        >>> dg.add_pair(2,5)
+        >>> dg.add_pair(5,6)
+        >>> dg.dfs(0,6)
+        [0, 2, 5, 6]
+        >>> dg.dfs(1,6)
+        [1, 5, 6]
+        >>> dg.dfs()
+        [0, 2, 5, 6, 1, 4, 3]
+        >>> dg.dfs(1,0)
+        [1, 5, 6, 4, 3]
+        """
         stack = []
         visited = []
         if s == -2:
-            s = next(iter(self.graph))
-        stack.append(s)
-        visited.append(s)
-        ss = s
-
-        while True:
-            # check if there is any non isolated nodes
-            if len(self.graph[s]) != 0:
-                ss = s
-                for node in self.graph[s]:
-                    if visited.count(node[1]) < 1:
-                        if node[1] == d:
-                            visited.append(d)
-                            return visited
-                        else:
-                            stack.append(node[1])
-                            visited.append(node[1])
-                            ss = node[1]
-                            break
-
-            # check if all the children are visited
-            if s == ss:
-                stack.pop()
-                if len(stack) != 0:
-                    s = stack[len(stack) - 1]
+            if self.graph.get(s,None):
+                pass # -2 is a node
+            elif len(self.graph) > 0:
+                s = next(iter(self.graph))
             else:
-                s = ss
+                return [] # Graph empty
+        stack.append(s)
 
-            # check if se have reached the starting point
-            if len(stack) == 0:
-                return visited
+        # Run dfs
+        while len(stack) > 0:
+            s = stack.pop()
+            visited.append(s)
+            # If reached d, return
+            if s==d:
+                break
+
+            # add not visited child nodes to stack
+            for _,ss in self.graph[s]:
+                if visited.count(ss) < 1:
+                    stack.append(ss)
+        return visited
 
     # c is the count of nodes you want and if you leave it or pass -1 to the function
     # the count will be random from 10 to 10000
@@ -84,12 +137,42 @@ def fill_graph_randomly(self, c=-1):
                     self.add_pair(i, n, 1)
 
     def bfs(self, s=-2):
+        """
+        Performs breadth first search from s
+        Returns list.
+        >>> dg = DirectedGraph()
+        >>> dg.bfs()
+        []
+        >>> dg.add_pair(1,1)
+        >>> dg.bfs(1)
+        [1]
+        >>> dg = DirectedGraph()
+        >>> dg.add_pair(0,1)
+        >>> dg.add_pair(0,2)
+        >>> dg.add_pair(1,3)
+        >>> dg.add_pair(1,4)
+        >>> dg.add_pair(1,5)
+        >>> dg.add_pair(2,5)
+        >>> dg.add_pair(5,6)
+        >>> dg.bfs(0)
+        [0, 1, 2, 3, 4, 5, 6]
+        >>> dg.bfs(1)
+        [1, 3, 4, 5, 6]
+        >>> dg.bfs()
+        [0, 1, 2, 3, 4, 5, 6]
+        """
         d = deque()
         visited = []
         if s == -2:
-            s = next(iter(self.graph))
+            if self.graph.get(s,None):
+                pass # -2 is a node
+            elif len(self.graph) > 0:
+                s = next(iter(self.graph))
+            else:
+                return [] # Graph empty
         d.append(s)
         visited.append(s)
+        # Run bfs
         while d:
             s = d.popleft()
             if len(self.graph[s]) != 0:
@@ -300,42 +383,60 @@ def remove_pair(self, u, v):
 
     # if no destination is meant the default value is -1
     def dfs(self, s=-2, d=-1):
-        if s == d:
-            return []
+        """
+        Performs depth first search from s to find d.
+        Returns the path s->d as a list.
+        Returns dfs from s if d is not found
+        >>> ug = Graph()
+        >>> ug.dfs()
+        []
+        >>> ug.add_pair(1,1)
+        >>> ug.dfs(1,1)
+        [1]
+        >>> ug = Graph()
+        >>> ug.add_pair(0,1)
+        >>> ug.add_pair(0,2)
+        >>> ug.add_pair(1,3)
+        >>> ug.add_pair(1,4)
+        >>> ug.add_pair(1,5)
+        >>> ug.add_pair(2,5)
+        >>> ug.add_pair(5,6)
+        >>> ug.dfs(0,6)
+        [0, 2, 5, 6]
+        >>> ug.dfs(1,6)
+        [1, 5, 6]
+        >>> ug.dfs()
+        [0, 2, 5, 6, 1, 4, 3]
+        >>> ug.dfs(1,0)
+        [1, 5, 6, 2, 0]
+        """
         stack = []
         visited = []
         if s == -2:
-            s = next(iter(self.graph))
+            if self.graph.get(s,None):
+                pass # -2 is a node
+            elif len(self.graph) > 0:
+                s = next(iter(self.graph))
+            else:
+                return [] # Graph empty
         stack.append(s)
-        visited.append(s)
-        ss = s
-
-        while True:
-            # check if there is any non isolated nodes
-            if len(self.graph[s]) != 0:
-                ss = s
-                for node in self.graph[s]:
-                    if visited.count(node[1]) < 1:
-                        if node[1] == d:
-                            visited.append(d)
-                            return visited
-                        else:
-                            stack.append(node[1])
-                            visited.append(node[1])
-                            ss = node[1]
-                            break
 
-            # check if all the children are visited
-            if s == ss:
-                stack.pop()
-                if len(stack) != 0:
-                    s = stack[len(stack) - 1]
+        # Run dfs
+        while len(stack) > 0:
+            s = stack.pop()
+            if visited.count(s) == 1:
+                continue
             else:
-                s = ss
-
-            # check if se have reached the starting point
-            if len(stack) == 0:
-                return visited
+                visited.append(s)
+            # If reached d, return
+            if s==d:
+                break
+
+            # add not visited child nodes to stack
+            for _,ss in self.graph[s]:
+                if visited.count(ss) < 1:
+                    stack.append(ss)
+        return visited
 
     # c is the count of nodes you want and if you leave it or pass -1 to the function
     # the count will be random from 10 to 10000
@@ -350,10 +451,39 @@ def fill_graph_randomly(self, c=-1):
                     self.add_pair(i, n, 1)
 
     def bfs(self, s=-2):
+        """
+        Performs breadth first search from s
+        Returns list.
+        >>> ug = Graph()
+        >>> ug.bfs()
+        []
+        >>> ug.add_pair(1,1)
+        >>> ug.bfs(1)
+        [1]
+        >>> ug = Graph()
+        >>> ug.add_pair(0,1)
+        >>> ug.add_pair(0,2)
+        >>> ug.add_pair(1,3)
+        >>> ug.add_pair(1,4)
+        >>> ug.add_pair(1,5)
+        >>> ug.add_pair(2,5)
+        >>> ug.add_pair(5,6)
+        >>> ug.bfs(0)
+        [0, 1, 2, 3, 4, 5, 6]
+        >>> ug.bfs(1)
+        [1, 0, 3, 4, 5, 2, 6]
+        >>> ug.bfs()
+        [0, 1, 2, 3, 4, 5, 6]
+        """
         d = deque()
         visited = []
         if s == -2:
-            s = next(iter(self.graph))
+            if self.graph.get(s,None):
+                pass # -2 is a node
+            elif len(self.graph) > 0:
+                s = next(iter(self.graph))
+            else:
+                return [] # Graph empty
         d.append(s)
         visited.append(s)
         while d: