邓心一 %!s(int64=6) %!d(string=hai) anos
pai
achega
085e88e9a2
Modificáronse 1 ficheiros con 17 adicións e 10 borrados
  1. 17 10
      oj/algorithms/graph/dijkstra.go

+ 17 - 10
oj/algorithms/graph/dijkstra.go

@@ -25,14 +25,19 @@ func main() {
 	fmt.Println(prev)
 }
 
-type minHeap [][2]int
+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][0] < h[j][0] }
+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.([2]int))
+	*h = append(*h, x.(pair))
 }
 
 func (h *minHeap) Pop() interface{} {
@@ -46,24 +51,26 @@ 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, [2]int{0, src})
+	heap.Push(&mh, pair{src, 0})
 	for mh.Len() != 0 {
-		u := heap.Pop(&mh).([2]int)
-		if dist[u[1]] < u[0] {
+		u := heap.Pop(&mh).(pair)
+		if used[u.k] {
 			continue
 		}
-		for _, v := range adj[u[1]] {
-			alt := u[0] + length[u[1]][v]
+		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[1]
-				heap.Push(&mh, [2]int{alt, v})
+				prev[v] = u.k
+				heap.Push(&mh, pair{v, alt})
 			}
 		}
 	}