examples/dijkstra.py
author Eugen Sawin <sawine@me73.com>
Tue, 22 Feb 2011 18:58:21 +0100
changeset 2 7b0f43733557
permissions -rw-r--r--
Added example hacks for some algorithms.
     1 """
     2 Name: dijkstra - Dijkstra for the Single Source Shortest Path problem
     3 
     4 Description: dijkstra applies the Dijkstra algorithm for the Single Source Shortest Path problem.
     5 
     6 Author: Eugen Sawin <sawine@me73.com>
     7 """
     8 
     9 def main():
    10     args = parse_arguments()
    11     s = args.source
    12     raw_graph = args.graph
    13     graph = [e.strip(' ,(').split(',') for e in raw_graph.split(')')]
    14     graph = [(e[0].strip(), e[1].strip(), int(e[2].strip())) for e in graph if len(e) == 3]
    15     distances = dijkstra(s, create_graph(graph))
    16     for (v, d) in distances.iteritems():
    17         print '%s: %i' % (v, d)
    18 
    19 def create_graph(desc):
    20     vertices = set([e[0] for e in desc]) 
    21     vertices = vertices.union([e[1] for e in desc])
    22     edges = dict(zip([(v1, v2) for (v1, v2, _) in desc], [c for (_, _, c) in desc]))
    23     return (vertices, edges)    
    24 
    25 from heapq import *
    26 
    27 def dijkstra(s, (vertices, edges)):
    28     inf = sum(edges.itervalues()) + 1
    29     dist = dict([(v, inf) for v in vertices])
    30     dist[s] = 0
    31     u = []
    32     for (v, d) in dist.iteritems():
    33         heappush(u, (d, v))
    34     active = vertices
    35     while len(u) and len(active):
    36         (_, node) = heappop(u)
    37         if node in active:
    38             active.remove(node)
    39             for (v1, v2) in [e for e in edges.iterkeys() if e[0] == node]:
    40                 c = edges[(v1, v2)]
    41                 if dist[v2] > dist[v1] + c:
    42                     dist[v2] = dist[v1] + c
    43                     heappush(u, (dist[v2], v2))
    44     return dist                    
    45 
    46 from argparse import ArgumentParser
    47 
    48 def parse_arguments():
    49     parser = ArgumentParser(description='Applies the Dijkstra algorithm for the single source shortest path problem.')
    50     parser.add_argument('graph', help='sequence of (u, v, c(u,v))')
    51     parser.add_argument('source', help='source vertex')
    52     return parser.parse_args()
    53 
    54 if __name__ == "__main__":
    55     main()