Python Program for Dijkstra's Shortest Path Algorithm

Finds the shortest distance from one node to every other in a weighted graph, always going on from the closest node, with heapq.

Dijkstra's algorithm finds the shortest distance from one node to every other in a graph whose edges have lengths, or weights: towns joined by roads of different lengths, say. It always goes on from the closest node it hasn't finished with, and a heap finds that node quickly. In the program's graph the direct road from A to B is 4 long, but going through C is shorter: 1 + 2 = 3. Every length must be 0 or more; graphs with negative weights need another algorithm, such as Bellman-Ford.

Program

import heapq


def dijkstra(graph, start):
    """Shortest distances from start to every reachable node."""
    distances = {start: 0}
    heap = [(0, start)]  # (distance so far, node), closest first
    while heap:
        distance, node = heapq.heappop(heap)
        if distance > distances[node]:
            continue  # out of date: a shorter way was found
        for neighbour, weight in graph[node]:
            new_distance = distance + weight
            if new_distance < distances.get(neighbour, float("inf")):
                distances[neighbour] = new_distance
                heapq.heappush(heap, (new_distance, neighbour))
    return distances


# Each node, and (node, length) for each one-way road leaving it
graph = {
    "A": [("B", 4), ("C", 1)],
    "B": [("D", 1)],
    "C": [("B", 2), ("D", 5)],
    "D": [("E", 3)],
    "E": [],
}
print("Distances from A:", dijkstra(graph, "A"))

Open in the editor Download dijkstra_shortest_path.py

Output

Distances from A: {'A': 0, 'B': 3, 'C': 1, 'D': 4, 'E': 7}

How it works

Visualize Step through the program line by line: each step below shows up as its line runs.

1. The heapq module keeps a list arranged so that its smallest item is always first. heappush() adds an item and heappop() takes the smallest, both quickly, however long the list.
2. distances holds the shortest distance found so far to each node: 0 for the start. The heap holds (distance, node) pairs still to look at, starting with the start node. Tuples compare by their first item, so the heap always gives the closest node first.
3. While the heap has entries, the closest one comes off. The first time a node comes off, its distance is final: any other way to it would pass through a node at least as far away, and no road has a negative length.
4. A node can be in the heap more than once, when a shorter way to it was found after it went in. An entry longer than the node's best distance is out of date, and is skipped.
5. For each road leaving the node, new_distance is the length of the way through this node. If that's shorter than any way found before (float("inf"), infinity, stands for no way yet), it becomes the neighbour's distance, and the neighbour goes into the heap with it.
6. When the heap is empty, distances holds the shortest distance to every node the start can reach.
7. The graph is a dictionary: each node's list holds a (neighbour, length) tuple for each road leaving it. The roads are one-way: C has a road to B, but B has none back to C.
8. The program prints the shortest distance from A to each node.

Each entry taken from the heap, with what's left in it afterwards, smallest first:

So A is 0, B 3 (through C), C 1, D 4 (through C and B) and E 7.

Time and space complexity

Each road can add one entry to the heap, so for V nodes and E roads the heap gets at most E + 1 entries, and each push and pop takes O(log E) steps, since a heap is a tree about log₂ E levels deep. A graph has fewer than V² edges, so log E is less than 2 log V, and the two are the same in Big-O. Each node's roads are read once, when its up-to-date entry comes off, so the total is O((V + E) log V). The heap can hold up to E entries and distances up to V, so the extra space is O(V + E).

Other ways to write it

Remembering the path

A second dict, previous, records the node each shortest distance came from. Following it back from the goal gives the route. This version also stops as soon as the goal comes off the heap, as its distance can't get any shorter after that.

import heapq

def shortest_path(graph, start, goal):
    distances = {start: 0}
    previous = {}
    heap = [(0, start)]
    while heap:
        distance, node = heapq.heappop(heap)
        if node == goal:
            break
        if distance > distances[node]:
            continue
        for neighbour, weight in graph[node]:
            new_distance = distance + weight
            if new_distance < distances.get(neighbour, float("inf")):
                distances[neighbour] = new_distance
                previous[neighbour] = node
                heapq.heappush(heap, (new_distance, neighbour))
    path = [goal]
    while path[-1] != start:
        path.append(previous[path[-1]])
    return path[::-1], distances[goal]

graph = {"A": [("B", 4), ("C", 1)], "B": [("D", 1)],
         "C": [("B", 2), ("D", 5)], "D": [("E", 3)], "E": []}
print(shortest_path(graph, "A", "E"))

Try it yourself

• Start from C. The result is {'C': 0, 'B': 2, 'D': 3, 'E': 6}, with no A: no road leads back to it.
• Make the road from C to B 5 long instead of 2. Going through C no longer pays, so B is 4 again, D 5 and E 8.
• Add a node F and a road to it from E, "E": [("F", 2)] and "F": []. F's distance is 9.
• Print each node and its distance on a line of its own, in alphabetical order with sorted().

More graphs programs

• Python Program for Breadth-First Search (BFS) Intermediate
• Python Program for Depth-First Search (DFS) Intermediate
• Python Program for Topological Sort (Kahn's Algorithm) Advanced
• Python Program for Union-Find (Disjoint Set) Advanced