Replace container/heap with native implementation.

Author: korthaj

heap.go

@@ -1,67 +1,109 @@
 package graph
 
-import (
-	"container/heap"
-)
-
 type prioQueue struct {
 	heap  []int // vertices in heap order
 	index []int // index of each vertex in the heap
 	cost  []int64
 }
 
-func emptyQueue(cost []int64) *prioQueue {
+func emptyPrioQueue(cost []int64) *prioQueue {
 	return &prioQueue{
 		index: make([]int, len(cost)),
 		cost:  cost,
 	}
 }
 
-func newQueue(cost []int64) *prioQueue {
+func newPrioQueue(cost []int64) *prioQueue {
 	n := len(cost)
-	h := &prioQueue{
+	q := &prioQueue{
 		heap:  make([]int, n),
 		index: make([]int, n),
 		cost:  cost,
 	}
-	for i := range h.heap {
-		h.heap[i] = i
-		h.index[i] = i
+	for i := range q.heap {
+		q.heap[i] = i
+		q.index[i] = i
 	}
-	return h
+	return q
 }
 
-func (m *prioQueue) Len() int { return len(m.heap) }
+// Len returns the number of elements in the queue.
+func (q *prioQueue) Len() int {
+	return len(q.heap)
+}
 
-func (m *prioQueue) Less(i, j int) bool {
-	return m.cost[m.heap[i]] < m.cost[m.heap[j]]
+// Push pushes v onto the queue.
+// The time complexity is O(log n) where n = q.Len().
+func (q *prioQueue) Push(v int) {
+	n := q.Len()
+	q.heap = append(q.heap, v)
+	q.index[v] = n
+	q.up(n)
 }
 
-func (m *prioQueue) Swap(i, j int) {
-	m.heap[i], m.heap[j] = m.heap[j], m.heap[i]
-	m.index[m.heap[i]] = i
-	m.index[m.heap[j]] = j
+// Pop removes the minimum element from the queue and returns it.
+// The time complexity is O(log n) where n = q.Len().
+func (q *prioQueue) Pop() int {
+	n := q.Len() - 1
+	q.swap(0, n)
+	q.down(0, n)
+
+	v := q.heap[n]
+	q.index[v] = -1
+	q.heap = q.heap[:n]
+	return v
 }
 
-func (pq *prioQueue) Push(x interface{}) {
-	n := len(pq.heap)
-	v := x.(int)
-	pq.heap = append(pq.heap, v)
-	pq.index[v] = n
+// Contains tells whether v is in the queue.
+func (q *prioQueue) Contains(v int) bool {
+	return q.index[v] >= 0
 }
 
-func (m *prioQueue) Pop() interface{} {
-	n := len(m.heap) - 1
-	v := m.heap[n]
-	m.index[v] = -1
-	m.heap = m.heap[:n]
-	return v
+// Fix re-establishes the ordering after the cost for v has changed.
+// The time complexity is O(log n) where n = q.Len().
+func (q *prioQueue) Fix(v int) {
+	if i := q.index[v]; !q.down(i, q.Len()) {
+		q.up(i)
+	}
 }
 
-func (m *prioQueue) Update(v int) {
-	heap.Fix(m, m.index[v])
+func (q *prioQueue) less(i, j int) bool {
+	return q.cost[q.heap[i]] < q.cost[q.heap[j]]
 }
 
-func (m *prioQueue) Contains(v int) bool {
-	return m.index[v] >= 0
+func (q *prioQueue) swap(i, j int) {
+	q.heap[i], q.heap[j] = q.heap[j], q.heap[i]
+	q.index[q.heap[i]] = i
+	q.index[q.heap[j]] = j
+}
+
+func (q *prioQueue) up(j int) {
+	for {
+		i := (j - 1) / 2 // parent
+		if i == j || !q.less(j, i) {
+			break
+		}
+		q.swap(i, j)
+		j = i
+	}
+}
+
+func (q *prioQueue) down(i0, n int) bool {
+	i := i0
+	for {
+		j1 := 2*i + 1
+		if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
+			break
+		}
+		j := j1 // left child
+		if j2 := j1 + 1; j2 < n && q.less(j2, j1) {
+			j = j2 // = 2*i + 2  // right child
+		}
+		if !q.less(j, i) {
+			break
+		}
+		q.swap(i, j)
+		i = j
+	}
+	return i > i0
 }

mst.go

@@ -1,9 +1,5 @@
 package graph
 
-import (
-	"container/heap"
-)
-
 // MST computes a minimum spanning tree for each connected component
 // of an undirected weighted graph.
 // The forest of spanning trees is returned as a slice of parent pointers:
@@ -22,13 +18,13 @@ func MST(g Iterator) (parent []int) {
 	}
 
 	// Prim's algorithm
-	Q := newQueue(cost)
+	Q := newPrioQueue(cost)
 	for Q.Len() > 0 {
-		v := heap.Pop(Q).(int)
+		v := Q.Pop()
 		g.Visit(v, func(w int, c int64) (skip bool) {
 			if Q.Contains(w) && c < cost[w] {
 				cost[w] = c
-				Q.Update(w)
+				Q.Fix(w)
 				parent[w] = v
 			}
 			return

path.go

@@ -1,9 +1,5 @@
 package graph
 
-import (
-	"container/heap"
-)
-
 // ShortestPath computes a shortest path from v to w.
 // Only edges with non-negative costs are included.
 // The number dist is the length of the path, or -1 if w cannot be reached.
@@ -44,10 +40,10 @@ func ShortestPaths(g Iterator, v int) (parent []int, dist []int64) {
 	dist[v] = 0
 
 	// Dijkstra's algorithm
-	Q := emptyQueue(dist)
-	heap.Push(Q, v)
+	Q := emptyPrioQueue(dist)
+	Q.Push(v)
 	for Q.Len() > 0 {
-		v = heap.Pop(Q).(int)
+		v := Q.Pop()
 		g.Visit(v, func(w int, d int64) (skip bool) {
 			if d < 0 {
 				return
@@ -56,10 +52,10 @@ func ShortestPaths(g Iterator, v int) (parent []int, dist []int64) {
 			switch {
 			case dist[w] == -1:
 				dist[w], parent[w] = alt, v
-				heap.Push(Q, w)
+				Q.Push(w)
 			case alt < dist[w]:
 				dist[w], parent[w] = alt, v
-				Q.Update(w)
+				Q.Fix(w)
 			}
 			return
 		})