Dijkstra

Share:
 1type edge struct{ to, co int }
 2type vert struct{ d, v int }
 3type pqForDijkstra []vert
 4
 5func (pq pqForDijkstra) Len() int            { return len(pq) }
 6func (pq pqForDijkstra) Swap(i, j int)       { pq[i], pq[j] = pq[j], pq[i] }
 7func (pq pqForDijkstra) Less(i, j int) bool  { return pq[i].d < pq[j].d }
 8func (pq *pqForDijkstra) Push(x interface{}) { *pq = append(*pq, x.(vert)) }
 9func (pq *pqForDijkstra) Pop() interface{} {
10	x := (*pq)[len(*pq)-1]
11	*pq = (*pq)[0 : len(*pq)-1]
12	return x
13}
14
15func dijkstra(es [][]edge, s, inf int) []int {
16	d := make([]int, len(es))
17	for i := 0; i < len(es); i++ {
18		d[i] = inf
19	}
20	d[s] = 0
21	pq := &pqForDijkstra{vert{0, s}}
22	heap.Init(pq)
23	for len(*pq) != 0 {
24		p := heap.Pop(pq).(vert)
25		if p.d > d[p.v] {
26			continue
27		}
28		for _, e := range es[p.v] {
29			cost := d[p.v] + e.co
30			if cost < d[e.to] {
31				d[e.to] = cost
32				heap.Push(pq, vert{d[e.to], e.to})
33			}
34		}
35	}
36	return d
37}
38