djikstra.py

from heapq import *
from collections import *

def djikstra(graph, start, end):
    """
    algo:
    - bfs with PQ
    - track total weight from start to a given vertex
      - initialize a table where each weight is infinite, except for the starting vertex (w = 0)
      - keep updating the total weight if we encounter a path that is smaller in weight
    - track prev vertex arriving at curr vertex
    """
    ## SETUP

    # unvisited in PQ sorted by distance
    unvisited, visited, distances, path = [], set(), {}, deque()
    distances[start] = (0, None)
    for u in graph.list.keys():
        if u == start:
            heappush(unvisited, (0, u))
        else:
            heappush(unvisited, (float('inf'), u))
    # NOTE: this can be skipped by checking for None before comparison
    for u in graph.list.keys():
        distances[u] = (float('inf'), None)

    ## ALGO IMPLEMENTATION
    while unvisited:
        (dist, u) = heappop(unvisited)
        if u not in visited:
            if u == end:
                break
            for vertex, vtx_dist in graph.list[u]:
                if vertex in visited:
                    continue
                new_distance = vtx_dist + dist
                if new_distance < distances[vertex][0]:
                    distances[vertex] = (new_distance, u)
                    heappush(unvisited, (new_distance, vertex))
            visited.add(u)

    ## PATH CONSTRUCTION
    # NOTE: we can optimize by storing the path in the priority queue
    vtx = end
    while vtx is not None:
        path.appendleft(vtx)
        d, src = distances[vtx]
        vtx = src
    return list(path)