one-file-projects/dijkstra.py

66 lines
1.4 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):
self.edges.add( (nodefrom,nodeto,extra) )
def all_edges_of(self,nodefrom):
return filter(lambda e: e[0] == nodefrom, self.edges)
def dijkstra(graph,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 main():
g = Graph()
for i in range(1,7):
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)
if __name__ == '__main__':
main()