Generic Priority Queue in Go

April 10, 2023

Go 1.18 introduced generics to the language, a highly requested feature. In this post we implement a priority queue in Go using generics to accept any element type in the queue without having to re-write boilerplate code.

Priority Queue

First, let's take a look at how we can actually implement a priority queue in Go. We can create a priority queue by implementing the heap interface that is present in the standard library (documentation):

// This example demonstrates a priority queue built using the heap interface.
package main

import (
	"container/heap"
	"fmt"
)

// An Item is something we manage in a priority queue.
type Item struct {
	value    string // The value of the item; arbitrary.
	priority int    // The priority of the item in the queue.
	// The index is needed by update and is maintained by the heap.Interface methods.
	index int // The index of the item in the heap.
}

// A PriorityQueue implements heap.Interface and holds Items.
type PriorityQueue []*Item

func (pq PriorityQueue) Len() int { return len(pq) }

func (pq PriorityQueue) Less(i, j int) bool {
	// We want Pop to give us the highest, not lowest, priority so we use greater than here.
	return pq[i].priority > pq[j].priority
}

func (pq PriorityQueue) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
	pq[i].index = i
	pq[j].index = j
}

func (pq *PriorityQueue) Push(x any) {
	n := len(*pq)
	item := x.(*Item)
	item.index = n
	*pq = append(*pq, item)
}

func (pq *PriorityQueue) Pop() any {
	old := *pq
	n := len(old)
	item := old[n-1]
	old[n-1] = nil  // avoid memory leak
	item.index = -1 // for safety
	*pq = old[0 : n-1]
	return item
}

As you can see, it is not exactly trivial to implement the heap interface. The heap interface is defined as such:

type Interface interface {
	sort.Interface
	Push(x any) // add x as element Len()
	Pop() any   // remove and return element Len() - 1.
}

And sort.Interface is defined as such:

type Interface interface {
	// Len is the number of elements in the collection.
	Len() int

	// Less reports whether the element with index i
	// must sort before the element with index j.
	//
	// If both Less(i, j) and Less(j, i) are false,
	// then the elements at index i and j are considered equal.
	// Sort may place equal elements in any order in the final result,
	// while Stable preserves the original input order of equal elements.
	//
	// Less must describe a transitive ordering:
	//  - if both Less(i, j) and Less(j, k) are true, then Less(i, k) must be true as well.
	//  - if both Less(i, j) and Less(j, k) are false, then Less(i, k) must be false as well.
	//
	// Note that floating-point comparison (the < operator on float32 or float64 values)
	// is not a transitive ordering when not-a-number (NaN) values are involved.
	// See Float64Slice.Less for a correct implementation for floating-point values.
	Less(i, j int) bool

	// Swap swaps the elements with indexes i and j.
	Swap(i, j int)
}

We have to implement a total of 5 functions just to create a priority queue for one custom type. That is a lot of boilerplate, especially if you need multiple priority queues for different custom types! This is where generics come in.

Generics

generic

/dʒəˈnerɪk/

shared by, including or typical of a whole group of things; not specific

One way to understand Generics is to think of it as something that enables type-agnostic operations. For example, in non-generic settings, when we define functions, we expect our parameters to be of a certain explicitly declared type, but with generics, as one might suspect from the term "generic", the parameter can be of any type*, and we can perform operations in these parameters without consideration for the parameter type. See this blog post for some quick examples and the official design document for more in-depth details.

In the case of our priority queue, we want to create a single priority queue structure that can accept any type of object in it. In order for our priority queue to work properly, we need to define certain things about the nature of our objects. In our priority queue, how do we know if one custom object is "prioritized" over another object?

* Subject to user defined constraints. See the "Constraints" section in the official design document for more information.

Natural Order

This leads us to the concept of Natural Order. The natural order of a domain of objects defines what is considered "smaller" or "bigger" for all objects in that domain. For example, we define the natural order of all natural numbers as 1 < 2 < 3 ... etc.

This allows us to also establish a priority for all objects in that domain; we can consider a smaller/bigger object in the natural order has having a bigger priority.

Comparator

How do we define the natural order of custom objects? We can use what is called a comparator function, which in Go is also just commonly called the "Less" function. A comparator function takes in two custom objects and computes and/or compares certain internal values, determining which object is "smaller". In a sense, this function is used to define the natural order. This is also the Less function that is defined in the heap Interface.

How the Comparator function is defined is up to your own object and use case; it can be as simple as comparing certain primitive values in our custom objects (this is probably the most common usage) or it can do certain pre-processing on various interal values of the object before comparing the final result.

Creating the Generic Priority Queue

Now we have all the ingredients needed for a generic priority queue. We need to implement the heap interface, and we also need a comparator function for the generic type which will be passed in by the user. The Generic Priority Queue is defined as such:

// Internally maintains a slice of pointers to items of type T.
type PQ[T comparable] struct {
	Items []*T
	less  LessFunction[T]
}

// Generic comparator function that is passed as a parameter when creating a new PQ.
type LessFunction[T comparable] func(T, T) bool

A priority queue of type T can be initialized with a comparator function for objects of type T:

// Creates a new instance of a priority queue with the given comparator function.
func CreatePQ[T comparable](less LessFunction[T]) *PQ[T] {
	return &PQ[T]{[]*T{}, less}
}

This will return a Priority Queue structure that will compare its items based on the comparator function we passed in.

Next, we want to implement the heap interface as defined by Go's standard library:

// Returns the size of the priority queue.
func (pq *PQ[T]) Len() int {
	return len(pq.Items)
}

// For implementing the heap interface; **do not use**.
func (pq *PQ[T]) Less(i, j int) bool {
	return pq.less(*pq.Items[i], *pq.Items[j])
}

// For implementing the heap interface; **do not use**.
func (pq *PQ[T]) Swap(i, j int) {
	pq.Items[i], pq.Items[j] = pq.Items[j], pq.Items[i]
}

// For implementing the heap interface; **do not use**.
func (pq *PQ[T]) Push(item any) {
	pq.Items = append(pq.Items, item.(*T))
}

func (pq *PQ[T]) Pop() any {
	old := pq.Items
	n := len(old)
	item := old[n-1]
	pq.Items = old[0 : n-1]
	return item
}

Note that these functions are only for implementing the heap interface, we should not be using them. By implementing these functions, we ensure that the first item in our PQ structure's internal slice of items is always of the highest priority.

With a non-generic priority queue, we would call heap.Push(pq, &item) and heap.Pop(pq) to perform Push and Pop operations. However, since we are creating a generic wrapper, we can encapsulate these in their own functions:

// Adds an item to the priority queue.
func (pq *PQ[T]) Put(item T) {
	heap.Push(pq, &item)
}

// Removes and returns the top item from the priority queue.
func (pq *PQ[T]) Get() T {
	return *heap.Pop(pq).(*T)
}

I have named them differently (Put and Get) in order to differentiate them from the functions implementing the heap interface. We can also add a Peek() function for completeness:

// Returns the top item from the priority queue without removing it.
func (pq *PQ[T]) Peek() T {
	return *pq.Items[0]
}

And with that, our generic priority queue is now usable:

type CustomItem struct{
	Key string
	Val int
}

// Create a PriorityQueue of type CustomItem by passing in a "less" comparator function to CreatePQ:
q := pq.CreatePQ[CustomItem](func(a, b CustomItem) bool { return a.Val < b.Val })

// Put items in the Priority Queue:
q.Put(CustomItem{"a", 5})
q.Put(CustomItem{"b", 1})

// Peek the head of the queue:
q.Peek() // CustomItem{"b", 1}

// Get (Dequeue) the head of the queue:
q.Get() // CustomItem{"b", 1}

This is definitely a much more succinct alternative to having to define all the functions ourselves to use the heap interface!

Closing Thoughts

Generics are extremely useful, especially if we know that we will be reusing similar data structures for different types of objects. The issue described here about the heap interface being too verbose to implement for different types was one I faced myself, and that led me to searching for answers for a more generic priority queue, which led me to discover Generics in Go (haha).

I'm definitely looking forward to using this new item in my toolbox in the future with similar problems!