Problem overview
The method dijkstra receives n, the number of vertices (labeled 0 to n-1), and edges, where each edge is [u, v, weight]. The graph is undirected and connected, and every weight is positive. The method returns an array dist in which dist[i] is the length of the shortest path from vertex 0 to vertex i.
For n = 5 and edges = [[0,1,4],[0,2,1],[2,1,2],[1,3,1],[2,3,5],[3,4,3]] the answer is [0, 3, 1, 4, 7]. The direct edge from 0 to 1 costs 4, but going through vertex 2 costs 1 + 2 = 3, so the longer-looking route wins.
Approach: always expand the closest unfinished vertex
Dijkstra's algorithm keeps a best-known distance for every vertex. It starts with 0 for the source and "unreached" for everyone else. It then repeats one move: take the unfinished vertex with the smallest known distance and relax its edges. Relaxing the edge from u to v with weight w means asking whether dist[u] + w beats the current dist[v], and updating it if so.
Why is the smallest vertex safe to finish? Every weight is at least 1, so any other route to it must pass through a vertex that is at least as far away. Nothing can shorten it later. That is why the algorithm needs non-negative weights.
The usual tool for "take the smallest" is a priority queue. Swift's standard library has none, so this lesson uses a plain array and sorts it before every removal. That costs more than a heap, but it is short and fast enough for the limits (at most 1000 vertices and 5000 edges).
Step by step
- Build adjacency lists.
adjis an array ofnlists of(neighbor, weight)tuples. Usingrepeating:count:with an empty array is safe because arrays are values: each slot gets its own independent list. - Add every edge twice. The graph is undirected, so the edge
[u, v, w]goes into the list ofuand into the list ofv. - Initialize distances. Every entry of
diststarts atInt.max, meaning "not reached yet", anddist[0]becomes 0. - Seed the queue.
pqholds(distance, vertex)pairs, distance first so sorting by the first element orders them by distance. It starts with(0, 0). - Loop while the queue has entries.
- Take the closest entry. Sort
pqascending by distance, then remove the first pair intodandu. Ifdis larger thandist[u], the pair is outdated and the loop skips it. - Visit each neighbor in
adj[u], unpacked asvandw. - Relax. If
dist[u] + wis smaller thandist[v], store the new distance and append the pair(dist[v], v)to the queue. - Return
dist. Once the queue is empty, every entry is final.
The solution in Swift
class Solution {
func dijkstra(_ n: Int, _ edges: [[Int]]) -> [Int] {
var adj = [[(Int, Int)]](repeating: [], count: n)
for e in edges {
adj[e[0]].append((e[1], e[2]))
adj[e[1]].append((e[0], e[2]))
}
var dist = [Int](repeating: Int.max, count: n)
dist[0] = 0
var pq = [(0, 0)]
while !pq.isEmpty {
pq.sort { $0.0 < $1.0 }
let (d, u) = pq.removeFirst()
if d > dist[u] {
continue
}
for (v, w) in adj[u] {
if dist[u] + w < dist[v] {
dist[v] = dist[u] + w
pq.append((dist[v], v))
}
}
}
return dist
}
}
The Swift details worth a closer look:
- Tuples in arrays.
[(Int, Int)]is an array of pairs. Inside the sort closure,$0.0is the first element of the first pair being compared and$1.0the first element of the second.let (d, u) = ...andfor (v, w) in ...unpack a pair into two names. - Type inference for the queue.
[(0, 0)]is inferred as[(Int, Int)], soappend((dist[v], v))needs the double parentheses: the inner ones build the tuple. Int.maxand overflow. Arithmetic overflow in Swift traps and ends the program, soInt.max + wwould crash. The code never computes it: only a vertex taken from the queue is expanded, and a vertex enters the queue only after receiving a real distance, sodist[u]is neverInt.maxwhendist[u] + wruns.- Lazy deletion. The code never changes a queue entry in place. When a shorter path appears, it adds a new pair and leaves the old one. In the example, vertex 1 is first queued with distance 4, then with distance 3. The pair
(4, 1)comes out later and is skipped by thed > dist[u]check. removeFirst(). It returns the removed element, and on an array it shifts all the remaining elements. That is another reason this queue is slower than a real heap.- Value semantics.
adj,distandpqarevararrays owned by this call, so mutating them never affects the caller's data.
Complexity
Let V be the number of vertices and E the number of edges.
Time. A successful relaxation adds one queue entry, and each directed edge relaxes successfully at most once per time its start vertex is expanded, which happens once. So the queue sees at most 2E + 1 entries in total. Each iteration sorts the queue and removes from its front, and the queue never holds more than 2E + 1 entries, so one iteration costs O(E log E), and the whole run is O(E² log E) in the worst case. With a binary heap the same algorithm runs in O(E log V). At 5000 edges the sorted array is still fine.
Space: O(V + E). The adjacency lists hold 2E entries, dist holds V, and the queue holds up to 2E + 1 pairs.
Edge cases and common mistakes
- Forgetting the second direction. Adding each edge to only one list treats the graph as directed. Any route that uses an edge from its second endpoint to its first would be missed.
- Parallel edges. With
edges = [[0,1,7],[1,0,3]]andn = 2, the result is[0, 3]. Relaxation keeps the cheaper edge automatically. - Skipping the stale-entry check. Without
d > dist[u], outdated pairs are expanded again. The answer stays correct, but the work grows. - Appending without an improvement test. Adding an entry on every neighbor visit fills the queue with useless pairs.
- Negative weights. The finishing argument above fails, so this algorithm can give wrong answers. The constraints rule them out here.
- Unreachable vertices. Their entries would stay at
Int.max. The lesson guarantees a connected graph, so this does not occur.
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],[1,0,3]]gives[0, 3].
