Dijkstra algorithm

dijkstra.py

import heapq
# graph is represented by adjacency list: List[List[pair]]
# s is the source vertex
def dijkstra(graph, s):
    # set is used to mark finalized vertices
    visited = set()
    # an array to keep the distance from s to this vertex.
    # initialize all distances as infinite, except s
    dist = [float('inf')] * len(graph)
    dist[s] = 0
    # priority queue containing (distance, vertex)
    min_heap = [(0, s)]

    while min_heap:
        # pop the vertex with the minimum distance
        _, u = heapq.heappop(min_heap)
        # skip if the vertex has already been visited
        if u in visited:
            continue
        visited.add(u)

        for v, weight in graph[u]:
            if v not in visited:
                # If there is shorted path from s to v through u.
                # s -> u -> v
                if dist[v] > (dist[u] + weight):
                    # Updating distance of v
                    dist[v] = dist[u] + weight
                    # insert to the queue
                    heapq.heappush(min_heap, (dist[v], v))

    return dist