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.