Added bfs and dfs

dijkstra.py

@@ -24,13 +24,28 @@ class Graph:
     def add_node(self,node):
         self.nodes.add(node)
 
-    def add_edge(self,nodefrom, nodeto, extra=None):
+    def add_edge(self,nodefrom, nodeto, extra=None, both=False):
+        self.require_node(nodefrom)
+        self.require_node(nodeto)
+
         self.edges.add( (nodefrom,nodeto,extra) )
 
+        if both:
+            self.add_edge(nodeto, nodefrom, extra=extra, both=False)
+
     def all_edges_of(self,nodefrom):
         return filter(lambda e: e[0] == nodefrom, self.edges)
 
+    def has_node(self, node):
+        return node in self.nodes
+
+    def require_node(self,node):
+        if not self.has_node(node):
+            raise Exception("Node %s does not exist" % repr(node))
+
+
 def dijkstra(graph,startnode):
+    graph.require_node(startnode)
     d = { node: Infinity() for node in graph.nodes }
     d[startnode] = 0
     V = [startnode]
@@ -45,9 +60,32 @@ def dijkstra(graph,startnode):
 
     return d
 
+def bfs(graph, startnode):
+    graph.require_node(startnode)
+    visited = set()
+    not_visited = [startnode]
+    while len(not_visited) > 0:
+        node = not_visited.pop(0)
+        yield node
+        visited.add(node)
+        for _, neighbor, extra in graph.all_edges_of(node):
+            if neighbor not in visited and neighbor not in not_visited:
+                not_visited.append(neighbor)
+
+def dfs(graph, node, visited=None):
+    graph.require_node(node)
+    if visited is None:
+        visited = []
+    yield node
+    visited.append(node)
+    for _, neighbor, extra in graph.all_edges_of(node):
+        if neighbor not in visited:
+            for node in dfs(graph, neighbor, visited=visited):
+                yield node
+
 def main():
     g = Graph()
-    for i in range(1,7):
+    for i in range(1,6):
         g.add_node(i)
 
     g.add_edge(1,2,10)
@@ -62,5 +100,21 @@ def main():
 
     print(d)
 
+    print( list(bfs(g,1)) )
+    print( list(dfs(g,1)) )
+
+    g2 = Graph()
+    for i in range(1,10):
+        g2.add_node(i)
+
+    for i in range(1,5):
+        g2.add_edge(i,2*i)
+        g2.add_edge(i,2*i+1)
+
+    print( list(bfs(g2,1)) )
+    print( list(dfs(g2,1)) )
+
+
+
 if __name__ == '__main__':
     main()