Dijkstra's Algorithm in Python

dijkstra.py

from heapq import heapify, heappop
from binascii import a2b_hex
from collections import defaultdict
from struct import unpack
from sys import argv, exit

def parse_byte(bstring):
    return unpack('>B', a2b_hex(bstring))[0]

def dijkstra(graph, source):
    inf = float('inf')
    dist = defaultdict(lambda: inf)
    previous = {}
    visited = set()

    dist[source] = 0
    graph_nodes = [(dist[n], n) for n in graph.keys()]
    heapify(graph_nodes)
    while graph_nodes:
        _, u = heappop(graph_nodes)
        visited.add(u)
        if dist[u] == inf:
            break

        for v in graph[u]:
            if v in visited:
                continue
            alt = dist[u] + graph[u][v]
            if alt < dist[v]:
                i = graph_nodes.index((dist[v], v))
                graph_nodes[i] = (alt, v)
                heapify(graph_nodes)

                dist[v] = alt
                previous[v] = u

    return dist, previous

if __name__ == "__main__":
    if len(argv) < 3:
        print 'Usage: dijkstra file-name start-node'
        exit(1)

    graph = defaultdict(dict)
    with open(argv[1]) as f:
        for line in f:
            if line.strip().startswith('#') or not line.strip():
                continue
            src, weight, dst = [parse_byte(col) for col in line.split()]
            graph[src][dst] = weight
            graph[dst][src] = weight

                    
    distances, previous =  dijkstra(graph, parse_byte(argv[2]))
    for node in distances:
        print 'distance from source to', node, ':', distances[node]

    print ''
    for node in previous:
        print previous[node], 'was previous to', node

graph

# example graph for dijkstra.py
# source weight destination
  01     a2     03  
  03     04     02  
  01     03     02  
  04     0f     01  
  05     02     02  
  05     03     01  
  02     05     03  
  03     01     04