dijkstra.py

from heapq import heappop, heappush

def dijkstra(g, s, f):
    h = [(0,s)]
    cost, prev = {}, {}
    cost[s], prev[s] = 0, None

    while len(h) > 0:
        _,v = heappop(h)
        if v == f:
            break
        for adj, adj_cost in g[v]:
            new_cost = cost[v] + adj_cost
            if new_cost < cost.get(adj, float('inf')):
                cost[adj], prev[adj] = new_cost, v
                heappush(h, (new_cost, adj))

    curr, path = f, [f]
    while curr in prev and prev[curr] != None:
        curr = prev[curr]
        path.append(curr)
    return path[::-1], cost.get(f, float('inf'))


if __name__ == '__main__':
    g = {'s':[('v',1),('w',4)],'v':[('w',2),('t',6)],'w':[('t',3)],'t':[]}
    print(dijkstra(g, 's', 't') == (['s', 'v', 'w', 't'], 6))
    print(dijkstra(g, 's', 'e') == (['e'], float('inf')))