Problem overview
The method dijkstra receives the number of vertices n, labeled 0 to n - 1, and a list of weighted edges, each a triple [u, v, weight]. The graph is undirected, so every edge can be used in both directions, and every weight is a positive integer. The method returns a list dist where dist[i] is the length of the shortest path from vertex 0 to vertex i, counting the sum of the weights. Every vertex can be reached.
In the example, n = 5 and edges = [[0, 1, 4], [0, 2, 1], [2, 1, 2], [1, 3, 1], [2, 3, 5], [3, 4, 3]]. The direct edge from 0 to 1 costs 4, but going 0 to 2 to 1 costs only 1 + 2 = 3. Vertex 3 is then reached through 1 for 3 + 1 = 4, which beats the route through 2 at 1 + 5 = 6. Vertex 4 adds 3 more, giving 7. The result is [0, 3, 1, 4, 7].
Approach: always expand the closest unfinished vertex
Breadth-first search counts edges, so it can't handle weights: a path with fewer edges may be longer. Dijkstra's algorithm fixes this with a greedy rule. Keep the best known distance to each vertex, and always process next the vertex with the smallest distance that is not yet final.
The rule works because the weights are positive. The closest unprocessed vertex can't be improved by a route through some other unprocessed vertex, since that route would first have to travel at least as far to reach the other vertex, and then add something more. So when a vertex is taken as the closest, its distance is final.
When a vertex is processed, the algorithm looks at its neighbors and tries relaxing each edge: if going through this vertex gives a neighbor a shorter distance than the one recorded, the record is updated.
The loop needs a fast way to pick the smallest distance. A priority queue does that, and Python provides one in the heapq module. It keeps a plain list arranged so that the smallest item can be removed quickly. The heap holds (distance, vertex) tuples. Tuples compare item by item, so the distance decides the order.
Heaps can't update an entry in place. The standard workaround is lazy deletion: when a vertex gets a better distance, push a new entry and leave the old one in the heap. When an outdated entry is popped later, its distance is larger than the recorded one, and it is skipped.
On the example, the pops go like this:
- Pop (0, 0). Relax: vertex 1 gets 4 and vertex 2 gets 1.
- Pop (1, 2). Relax: vertex 1 improves to 3, and vertex 3 gets 6.
- Pop (3, 1). Vertex 3 improves to 4.
- Pop (4, 1), which is outdated, because vertex 1 is at 3. It is skipped.
- Pop (4, 3). Vertex 4 gets 7.
- Pop (6, 3), outdated, skipped. Pop (7, 4), which relaxes nothing.
Step by step
- Build the adjacency list.
adjis a list ofnseparate empty lists, one per vertex, holding(neighbor, weight)pairs. - Fill it by unpacking each edge into
u, v, w. Append(v, w)toadj[u]and(u, w)toadj[v], because the graph is undirected. - Create
distwithfloat('inf')in every slot to mean not reached, then setdist[0] = 0. - Create the heap
pqas the list[(0, 0)]: distance first, vertex second. A list with one item is already a valid heap. - Loop while the heap is not empty.
- Pop the smallest entry with
heapq.heappop(pq)and unpack it intodandu. - Skip outdated entries. If
d > dist[u], a shorter path touwas found after this entry was pushed, socontinue. - Loop over the neighbors of
u, unpacking each pair intovandw. - Relax the edge. If
dist[u] + w < dist[v], setdist[v]to the smaller value and push(dist[v], v)withheapq.heappush. Pushing only on a real improvement keeps the heap small. - Return
dist. When the heap is empty, every distance is final.
The solution in Python
class Solution:
def dijkstra(self, n: int, edges: List[List[int]]) -> List[int]:
adj = [[] for _ in range(n)]
for u, v, w in edges:
adj[u].append((v, w))
adj[v].append((u, w))
dist = [float('inf')] * n
dist[0] = 0
pq = [(0, 0)]
while pq:
d, u = heapq.heappop(pq)
if d > dist[u]:
continue
for v, w in adj[u]:
if dist[u] + w < dist[v]:
dist[v] = dist[u] + w
heapq.heappush(pq, (dist[v], v))
return dist
The Python details worth a closer look:
heapqis a min-heap on a plain list.heappopalways returns the smallest item, andheappushadds one while keeping the heap order. The list itself is not sorted, so don't readpq[1]and expect the second smallest. Onlypq[0]is guaranteed to be the smallest.- Tuple comparison.
(4, 1) < (4, 3)compares the first items, and if they tie, the second. Entries with equal distance are ordered by vertex number, which is fine because vertices are integers and can always be compared. float('inf'). It compares greater than every integer, so any real path length beats it. The list then mixes integers and one float value, which works for comparisons. For the guaranteed connected graph every entry ends as an integer.- Safe repetition.
[float('inf')] * nrepeats an immutable value, so the slots are independent. The adjacency list uses a comprehension because lists are mutable. - Unpacking in loop headers.
for u, v, w in edgesandfor v, w in adj[u]split each item into names, which reads better than indexing withedge[0]. - Import. The code calls
heapq.heappopby its module name, so the file needsimport heapqat the top.
Complexity
Time: O((n + m) log n) for a graph with m edges, often written O(m log n) for a connected graph. Each edge can cause at most one push per direction, so the heap sees at most 2m pushes, plus the initial one. Every push and pop costs O(log) of the heap size, and that size is at most about 2m, whose logarithm is within a constant factor of log n because m is at most n².
Space: O(n + m). The adjacency list stores two entries per edge, dist has n entries, and the heap holds at most one entry per successful relaxation.
Edge cases and common mistakes
- A longer edge with fewer steps. For edges
[[0, 1, 1], [1, 2, 1], [0, 2, 5]]the distance to vertex 2 is 2, through vertex 1, not the direct 5. - Parallel edges. With edges
[[0, 1, 5], [0, 1, 2], [1, 2, 1]]the result is[0, 2, 3]: relaxation keeps the cheaper edge. - Two vertices.
n = 2with[[0, 1, 7]]gives[0, 7]. - Skipping the outdated-entry check. The result stays correct, but outdated entries then re-scan their neighbor lists and slow the algorithm down.
- Adding only one direction. The graph would be directed, and shortest paths through an edge used backwards would be missed.
- Negative weights. The greedy rule needs weights that are not negative, because a negative edge could improve a vertex that was already final. Here every weight is at least 1.
- Heap entries in the wrong order. Pushing
(v, dist)would order the heap by vertex number and lose the "closest first" rule. - Using a sorted list or a scan. Finding the minimum by scanning all vertices each time costs
O(n)per step, which isO(n²)overall.
Testing
n = 5,edges = [[0, 1, 4], [0, 2, 1], [2, 1, 2], [1, 3, 1], [2, 3, 5], [3, 4, 3]]gives[0, 3, 1, 4, 7].n = 2,edges = [[0, 1, 7]]gives[0, 7].n = 4,edges = [[0, 1, 1], [1, 2, 1], [0, 2, 5], [2, 3, 1]]gives[0, 1, 2, 3].n = 3,edges = [[0, 1, 5], [0, 1, 2], [1, 2, 1]]gives[0, 2, 3].