#!/usr/bin/env python class Infinity: def __eq__(self,other): return isinstance(other,Infinity) def __gt__(self,other): return not isinstance(other, Infinity) def __lt__(self,other): return False def __add__(self,other): return self def __radd__(self,other): return self def __repr__(self): return "inf" class Graph: def __init__(self): self.nodes = set() self.edges = set() def add_node(self,node): self.nodes.add(node) 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] while len(V) > 0: node = min(V, key=lambda e: d[e]) V.remove(node) for _, neighbor, distance in graph.all_edges_of(node): if d[neighbor] > d[node] + distance: d[neighbor] = d[node] + distance V.append(neighbor) 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,6): g.add_node(i) g.add_edge(1,2,10) g.add_edge(1,4,30) g.add_edge(1,5,100) g.add_edge(2,3,50) g.add_edge(3,5,10) g.add_edge(4,3,20) g.add_edge(4,5,60) d = dijkstra(g,1) 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()