167 lines
3.8 KiB
Python
167 lines
3.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 prim(graph,r=None):
|
|
if r is None:
|
|
r = list(graph.nodes)[0]
|
|
Vt = set();
|
|
Vt.add(r);
|
|
d = { node: Infinity() for node in graph.nodes }
|
|
d[r] = 0;
|
|
for _, neighbor, distance in graph.all_edges_of(r):
|
|
d[neighbor] = distance
|
|
|
|
while Vt < graph.nodes:
|
|
u = min(graph.nodes - Vt, key=lambda node: d[node])
|
|
if d[u] == Infinity():
|
|
break
|
|
Vt.add(u)
|
|
for _, neighbor, distance in graph.all_edges_of(u):
|
|
if neighbor not in Vt:
|
|
d[neighbor] = min([ d[neighbor], distance ])
|
|
|
|
return d;
|
|
|
|
|
|
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)
|
|
|
|
d = prim(g)
|
|
|
|
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)) )
|
|
|
|
g3 = Graph()
|
|
g3.add_node(1)
|
|
g3.add_node(2)
|
|
g3.add_node(3)
|
|
g3.add_node(4)
|
|
g3.add_node(5)
|
|
g3.add_node(6)
|
|
|
|
g3.add_edge(1,2,20)
|
|
g3.add_edge(1,3,10)
|
|
g3.add_edge(2,4,30)
|
|
g3.add_edge(2,5,20)
|
|
g3.add_edge(3,4,10)
|
|
g3.add_edge(3,5,20)
|
|
g3.add_edge(4,6,10)
|
|
g3.add_edge(5,6,30)
|
|
|
|
d = prim(g3,1)
|
|
print(d)
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|