dijkstra.py 文件源码

def dijkstra(graph, source):
    pq = []
    nodes = graph.get_all_vertices()
    distances = [math.inf] * len(nodes)
    path = [-1] * len(nodes)
    distances[source.index] = 0
    for node in nodes:
        # Store as (priority, task) tuples, heapq will sort on first element.
        heapq.heappush(pq, (distances[node.index], node))
    while pq:
        # Assumes non negative weights, so when popping a node it is the best way to get there.
        dist, node = heapq.heappop(pq)
        for edge in graph.get_adjacent_edges(node):
            # Note: can't terminate early and do this.
            # Eg: (s) -3-> (c) -12-> (d)
            #      \-20->(d) will be wrong
            # if distances[edge.destination.index] != math.inf:  # We already have the shortest path to this node.
            #     continue
            if relax(node, edge.destination, edge, distances):
                # Found a better way to get to a next node, add that to the pq and set the parent.
                heapq.heappush(pq, (distances[edge.destination.index], edge.destination))
                path[edge.destination.index] = node.index
    return distances, path  # Shortest path from source to any other node in distances.