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