leetcode-go/basis/algorithms/graph/dijkstra.go

package main

import (
	"container/heap"
	"fmt"
	"math"
)

func main() {
	var n, m, src int
	fmt.Scan(&n, &m, &src)
	adj := make([][]int, n)
	length := make([][]int, n)
	for i := range length {
		length[i] = make([]int, n)
	}
	for i := 0; i < m; i++ {
		var u, v, l int
		fmt.Scan(&u, &v, &l)
		adj[u] = append(adj[u], v)
		length[u][v] = l
	}
	dist, prev := dijkstra(adj, length, src)
	fmt.Println(dist)
	fmt.Println(prev)
}

type pair struct {
	k int
	w int
}

type minHeap []pair

func (h minHeap) Len() int           { return len(h) }
func (h minHeap) Less(i, j int) bool { return h[i].w < h[j].w }
func (h minHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *minHeap) Push(x interface{}) {
	*h = append(*h, x.(pair))
}

func (h *minHeap) Pop() interface{} {
	l := len(*h)
	x := (*h)[l-1]
	*h = (*h)[:l-1]
	return x
}

func dijkstra(adj [][]int, length [][]int, src int) (dist []int, prev []int) {
	n := len(adj)
	dist = make([]int, n)
	prev = make([]int, n)
	used := make([]bool, n)
	for i := 0; i < n; i++ {
		dist[i] = math.MaxInt32
		prev[i] = -1
	}
	dist[src] = 0
	var mh minHeap
	heap.Push(&mh, pair{src, 0})
	for mh.Len() != 0 {
		u := heap.Pop(&mh).(pair)
		if used[u.k] {
			continue
		}
		used[u.k] = true
		for _, v := range adj[u.k] {
			alt := u.w + length[u.k][v]
			if alt < dist[v] {
				dist[v] = alt
				prev[v] = u.k
				heap.Push(&mh, pair{v, alt})
			}
		}
	}
	return
}