examples/dijkstra.py
changeset 2 7b0f43733557
     1.1 --- /dev/null	Thu Jan 01 00:00:00 1970 +0000
     1.2 +++ b/examples/dijkstra.py	Tue Feb 22 18:58:21 2011 +0100
     1.3 @@ -0,0 +1,55 @@
     1.4 +"""
     1.5 +Name: dijkstra - Dijkstra for the Single Source Shortest Path problem
     1.6 +
     1.7 +Description: dijkstra applies the Dijkstra algorithm for the Single Source Shortest Path problem.
     1.8 +
     1.9 +Author: Eugen Sawin <sawine@me73.com>
    1.10 +"""
    1.11 +
    1.12 +def main():
    1.13 +    args = parse_arguments()
    1.14 +    s = args.source
    1.15 +    raw_graph = args.graph
    1.16 +    graph = [e.strip(' ,(').split(',') for e in raw_graph.split(')')]
    1.17 +    graph = [(e[0].strip(), e[1].strip(), int(e[2].strip())) for e in graph if len(e) == 3]
    1.18 +    distances = dijkstra(s, create_graph(graph))
    1.19 +    for (v, d) in distances.iteritems():
    1.20 +        print '%s: %i' % (v, d)
    1.21 +
    1.22 +def create_graph(desc):
    1.23 +    vertices = set([e[0] for e in desc]) 
    1.24 +    vertices = vertices.union([e[1] for e in desc])
    1.25 +    edges = dict(zip([(v1, v2) for (v1, v2, _) in desc], [c for (_, _, c) in desc]))
    1.26 +    return (vertices, edges)    
    1.27 +
    1.28 +from heapq import *
    1.29 +
    1.30 +def dijkstra(s, (vertices, edges)):
    1.31 +    inf = sum(edges.itervalues()) + 1
    1.32 +    dist = dict([(v, inf) for v in vertices])
    1.33 +    dist[s] = 0
    1.34 +    u = []
    1.35 +    for (v, d) in dist.iteritems():
    1.36 +        heappush(u, (d, v))
    1.37 +    active = vertices
    1.38 +    while len(u) and len(active):
    1.39 +        (_, node) = heappop(u)
    1.40 +        if node in active:
    1.41 +            active.remove(node)
    1.42 +            for (v1, v2) in [e for e in edges.iterkeys() if e[0] == node]:
    1.43 +                c = edges[(v1, v2)]
    1.44 +                if dist[v2] > dist[v1] + c:
    1.45 +                    dist[v2] = dist[v1] + c
    1.46 +                    heappush(u, (dist[v2], v2))
    1.47 +    return dist                    
    1.48 +
    1.49 +from argparse import ArgumentParser
    1.50 +
    1.51 +def parse_arguments():
    1.52 +    parser = ArgumentParser(description='Applies the Dijkstra algorithm for the single source shortest path problem.')
    1.53 +    parser.add_argument('graph', help='sequence of (u, v, c(u,v))')
    1.54 +    parser.add_argument('source', help='source vertex')
    1.55 +    return parser.parse_args()
    1.56 +
    1.57 +if __name__ == "__main__":
    1.58 +    main()