update

oj/algorithms/graph/dijkstra.go

@@ -0,0 +1,71 @@
+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 minHeap [][2]int
+
+func (h minHeap) Len() int           { return len(h) }
+func (h minHeap) Less(i, j int) bool { return h[i][0] < h[j][0] }
+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.([2]int))
+}
+
+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)
+	for i := 0; i < n; i++ {
+		dist[i] = math.MaxInt32
+		prev[i] = -1
+	}
+	dist[src] = 0
+	var mh minHeap
+	heap.Push(&mh, [2]int{0, src})
+	for mh.Len() != 0 {
+		u := heap.Pop(&mh).([2]int)
+		if dist[u[1]] < u[0] {
+			continue
+		}
+		for _, v := range adj[u[1]] {
+			alt := u[0] + length[u[1]][v]
+			if alt < dist[v] {
+				dist[v] = alt
+				prev[v] = u[1]
+				heap.Push(&mh, [2]int{alt, v})
+			}
+		}
+	}
+	return
+}