python implementation of dijkstra shortest path.

dijkstra.py

# -*- coding: utf-8 -*-


def dijkstra(graph, source=0):
    nodes = dict.fromkeys(range(len(graph)), 0)

    # all visited nodes, source is the first entry.
    visited = {source: 0}

    # previous node dict
    previous = {}

    # initial distance dict.
    p_dist = {source: 0}
    for i in nodes:
        p_dist[i] = graph[source][i]

    while set(visited) != set(nodes):
        for x in set(nodes) - set(visited):
            minx = x
        for x, val in p_dist.items():
            if x not in visited and val < p_dist[minx]:
                minx = x
        if graph[source][minx] != inf:
            previous[minx] = source
        visited[minx] = 0
        for y in set(nodes) - set(visited):
            if p_dist[y] > p_dist[minx] + graph[minx][y]:
                previous[y] = minx
                p_dist[y] = p_dist[minx] + graph[minx][y]

    return p_dist, previous

if __name__ == '__main__':
    inf = float('inf')

    graph = [
        [0, 2, 1, inf, inf, inf],
        [inf, 0, 3, 3, inf, inf],
        [inf, inf, 0, inf, 1, inf],
        [inf, inf, inf, 0, inf, 2],
        [inf, inf, inf, 2, 0, 5],
        [inf, inf, inf, inf, inf, 0],
    ]

    source = 0

    print dijkstra(graph, source)