one-file-projects/graph.py

121 lines
2.8 KiB
Python

#!/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()