diff --git a/dijkstra.py b/dijkstra.py index bdf7d9a..06bc21d 100644 --- a/dijkstra.py +++ b/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()