Using the go heap interface
07 March 2021
Go provides support for heaps as part of the language. What we have are interface and module methods, these can be a little intimidating if you were looking for a generic heap object out of the box.
Expectations
If you are expecting an object your API might look something like this, where pq
is shorthand for `priorityQueue:
pq := heap.NewHeap()
pq.Push(item)
pq.Pop()
But instead, we get module functions that take a heap interface and its arguments:
heap.Push(pq, item)
heap.Pop(pq)
Interface
This is the go heap interface:
type Interface interface {
sort.Interface
Push(x interface{})
Pop() interface{}
}
Plus the go sort interface:
type Interface interface {
Len() int
Less(i, j int) bool
Swap(i, j int)
}
Implementation
Putting them together we get a full implementation:
package main
import (
"container/heap"
"fmt"
)
type Item struct {
value int
priority int
}
type PriorityQueue []*Item
func (pq PriorityQueue) Len() int { return len(pq) }
func (pq PriorityQueue) Swap(i, j int) { pq[i], pq[j] = pq[j], pq[i] }
func (pq PriorityQueue) Less(i, j int) bool {
return pq[i].priority > pq[j].priority
}
func (pq *PriorityQueue) Push(x interface{}) {
// a failed type assertion here will panic
*pq = append(*pq, x.(*Item))
}
func (pq *PriorityQueue) Pop() interface{} {
old := *pq
n := len(old)
item := old[n-1]
old[n-1] = nil
*pq = old[0 : n-1]
return item
}
func main() {
pq := &PriorityQueue{}
heap.Push(pq, &Item{value: 100, priority: 10})
item := heap.Pop(pq)
fmt.Println(item)
}
Implementing this interface gives us the heap methods we expect:
- heap.Init(pq)
- heap.Push(pq, interface {})
- heap.Pop(pq)
As well as two helper methods, if you are tracking the item index in the underlying heap array:
- heap.Remove(pq, index), removing element at index
- heap.Fix(pq, index), equivalent to remove/push element at index
The downsides to the interface approach, besides having to implement five functions, are the interface type assertions from pop and push. The caller of pop needs to type assert, which is only a few lines of code. But passing the wrong type to push could lead to a panic, as it does in this example.
Without generics, this still provides the heap logic managed by the module.