LESSON 332 OF 338 · GRAPHS

Dijkstra's Algorithm in Go: Shortest Distances with sort.Slice

Practice Priority Queue, O((V+E) log V) Read the full explanation

Given a weighted undirected graph with n vertices (labeled 0 to n-1) and a list of edges where edges[i] = [u, v, weight], find the shortest distance from vertex 0 to every other vertex using Dijkstra's algorithm.

Return an array dist of length n where dist[i] is the shortest distance from vertex 0 to vertex i. All vertices are reachable.

Example

Input: n = 5, edges = [[0,1,4],[0,2,1],[2,1,2],[1,3,1],[2,3,5],[3,4,3]]
Output: [0, 3, 1, 4, 7]
Explanation:
- 0 → 0: distance 0
- 0 → 2 → 1: distance 1 + 2 = 3
- 0 → 2: distance 1
- 0 → 2 → 1 → 3: distance 1 + 2 + 1 = 4
- 0 → 2 → 1 → 3 → 4: distance 1 + 2 + 1 + 3 = 7

Constraints

2 <= n <= 1000
1 <= edges.length <= 5000
0 <= u, v < n
1 <= weight <= 1000
The graph is connected.

STEP 1 OF 9

Guided

Ctrl + Enter to check

GO · GRAPHS

Dijkstra in Go, explained

Problem overview

The function dijkstra receives the number of vertices n (labeled 0 to n - 1) and a list of edges, where each edge [u, v, weight] joins two vertices in both directions with a positive cost. It returns a slice dist where dist[i] is the length of the cheapest path from vertex 0 to vertex i. The graph is connected, so every vertex is reachable.

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 only 1 + 2 = 3. Vertex 3 is then 3 + 1 = 4 away, and vertex 4 is 4 + 3 = 7 away.

Approach: always settle the closest unsettled vertex

Breadth-first search counts edges, which is wrong when edges have different costs. Dijkstra's algorithm handles positive weights with a similar loop, but it keeps a priority queue ordered by distance so far instead of a first-in, first-out queue.

The idea: keep a best-known distance for every vertex, with infinity as the starting value, and 0 for the source. Repeatedly take the queued vertex with the smallest distance. Since all weights are positive, nothing found later can beat that distance, so it is final. Then relax each of its edges: if going through this vertex gives a neighbour a shorter distance than it has now, record the better one and queue the neighbour again.

Updating an entry inside a queue is awkward, so the queue simply receives a new entry when a shorter path is found. The older, worse entry stays and is skipped when it comes to the front. This is called lazy deletion: an entry is stale when its distance is larger than the vertex's current best distance.

Tracing the example: the queue starts as (0, vertex 0). Relaxing vertex 0 gives vertex 1 a distance of 4 and vertex 2 a distance of 1. Vertex 2 is closest, so it goes next and improves vertex 1 from 4 to 3, and gives vertex 3 a distance of 6. Vertex 1 at 3 is next, and it improves vertex 3 from 6 to 4. The entries (4, vertex 1) and (6, vertex 3) are now stale, and are skipped when they surface. Vertex 3 at 4 gives vertex 4 a distance of 7.

This lesson keeps the queue as a plain slice and sorts it before each pop. That is easy to write and fast enough for the lesson's limits. The standard approach uses a binary heap, which gets the smallest entry in O(log n) without re-sorting everything.

Step by step

Declare adj as a [][][2]int of length n. Each entry is a list of [2]int pairs holding a neighbour and a weight.
Add every edge to both endpoints' lists, because the graph is undirected. Assign the result of append back.
Create dist, fill it with math.MaxInt32 to mean "not reached", and set dist[0] = 0.
Create the queue pq as a [][2]int holding {0, 0}, the pair (distance, vertex), with the distance first so it is the sort key.
Loop while len(pq) > 0.
Sort, pop, skip stale entries. Sort pq by distance with sort.Slice, read the front into d and u, drop it by reslicing, and continue if d is larger than dist[u].
Range over adj[u], naming each entry nb.
Relax. Unpack nb into v and w. If dist[u]+w < dist[v], update dist[v] and append [2]int{dist[v], v} to the queue. Adding entries only on a real improvement keeps the queue small.
Return dist. When the queue is empty, every distance is final.

The solution in Go

func dijkstra(n int, edges [][]int) []int {
    adj := make([][][2]int, n)
    for _, e := range edges {
        adj[e[0]] = append(adj[e[0]], [2]int{e[1], e[2]})
        adj[e[1]] = append(adj[e[1]], [2]int{e[0], e[2]})
    }
    dist := make([]int, n)
    for i := range dist {
        dist[i] = math.MaxInt32
    }
    dist[0] = 0
    pq := [][2]int{{0, 0}}
    for len(pq) > 0 {
        sort.Slice(pq, func(i, j int) bool { return pq[i][0] < pq[j][0] })
        d, u := pq[0][0], pq[0][1]
        pq = pq[1:]
        if d > dist[u] {
            continue
        }
        for _, nb := range adj[u] {
            v, w := nb[0], nb[1]
            if dist[u]+w < dist[v] {
                dist[v] = dist[u] + w
                pq = append(pq, [2]int{dist[v], v})
            }
        }
    }
    return dist
}

The Go details worth knowing:

Imports. The code uses sort and math, so the program needs import "sort" and import "math".
sort.Slice takes the slice and a function func(i, j int) bool that reports whether element i must come before element j. It receives indices, so the function reads pq[i][0] and pq[j][0] from the slice. It sorts in place, and the sort is not stable, so equal distances may come out in any order. That doesn't matter here.
math.MaxInt32 is 2147483647. It is used as "infinity" and is far above any real distance. Because only reached vertices are expanded, dist[u] is never the sentinel when dist[u]+w is computed, and int is 64 bits here anyway, so the sum can't overflow.
Pairs as [2]int. A fixed-size array is a value, so each queue entry is a self-contained pair, and pq[0][0], pq[0][1] read its two parts.
Reslicing pops. pq = pq[1:] drops the first entry without copying. The next append may reallocate.
continue skips the rest of the loop body for a stale entry and starts the next pass.

Complexity

With a binary heap the time is O(E log V) for E edges and V vertices, because each successful relaxation adds one heap entry. This lesson's version sorts the whole queue on every pass, which costs O(P log P) for a queue of P entries. Each edge is relaxed at most twice in total (once from each end), so there are at most about 2E entries and pops, and the whole run is roughly O(E² log E) in the worst case. For up to 5000 edges that is acceptable, but a real heap scales far better. Space is O(V + E) for the adjacency lists, dist and the queue.

Edge cases and common mistakes

Parallel edges. With edges [[0,1,2],[0,1,1],[1,2,1]] the cheaper edge wins and the answer is [0, 1, 2].
A cheaper detour. In [[0,1,1],[1,2,1],[0,2,5]] vertex 2 is first reached at cost 5 and then improved to 2 through vertex 1. The stale entry (5, vertex 2) is skipped later.
Two vertices. n = 2 with the edge [0,1,7] returns [0, 7].
Forgetting the stale check. The result is still correct, but old entries get expanded again and the work grows.
Adding each edge in one direction only. The graph is undirected, so a path may need to cross an edge backwards.
Negative weights. The "first pop is final" argument needs non-negative weights, and the constraints here guarantee positive ones.
Disconnected graphs. This lesson guarantees connectivity. If a vertex were unreachable, its entry would stay at math.MaxInt32. For example, n = 3 with only the edge [0,1,1] gives [0 1 2147483647].

Testing

n = 5 with the example edges returns [0 3 1 4 7].
n = 3, edges [[0,1,1],[1,2,1],[0,2,5]] returns [0 1 2].
n = 2, edges [[0,1,7]] returns [0 7].

Keep going with GoAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.

Start learning free