In this post we will learn about heap data structures, what they are used for, and how to implement them in Go using the container/heap standard library
If you just want to see the example code, you can view it on Github or run the live example on the Go playground
What Is a Heap Data Structure?
A heap is a tree-based data structure with a fixed relationship between each parent node and its child nodes.
For max heaps, this means that the value of any parent node must be greater than or equal to its child node. This means that the root node will always have the maximum value.
Min heaps have the reverse property, an so the root node of a min heap will always have the least value.
Why Do We Use Heaps?
Normally, we use heaps by “heapifying” a set of values - which means converting a list of items into a heap:
Some of the key properties which make heaps useful are:
- We can find the maximum or minimum value of a given set of “n” values in O(1) time. If we use an array based list, this would normally take O(n) time.
- Adding and removing values takes O(log(n)) time, while preserving the max or min heap property. This is used when sorting a list of items using heap sort, where each element is inserted into a heap one at a time.
Using the Heap Interface in Go
We can convert a regular list of items into a heap by implementing the heap interface for that type.
The heap interface has five methods and once we implement them we can use the functions provided by the container/heap library.
For example let’s see how we can implement the heap interface for a list of integers:
// we need to define a custom type instead of using the raw integer slice
// since we need to define methods on the type to implement the heap interface
type intHeap []int
// Len is the number of elements in the collection.
func (h intHeap) Len() int {
return len(h)
}
// Less reports whether the element with index i
// must sort before the element with index j.
// This will determine whether the heap is a min heap or a max heap
func (h intHeap) Less(i int, j int) bool {
return h[i] < h[j]
}
// Swap swaps the elements with indexes i and j.
func (h intHeap) Swap(i int, j int) {
h[i], h[j] = h[j], h[i]
}
// Push and Pop are used to append and remove the last element of the slice
func (h *intHeap) Push(x any) {
*h = append(*h, x.(int))
}
func (h *intHeap) Pop() any {
old := *h
n := len(old)
x := old[n-1]
*h = old[0 : n-1]
return x
}
Now that our intHeap type implements the heap.Interface interface, we can use the functions provided by the container/heap package to perform various operations.
For example, let’s start by heapifying an intHeap instance:
package main
import (
"container/heap"
"fmt"
)
func main() {
// create a new intHeap instance
nums := &intHeap{3, 1, 4, 5, 1, 1, 2, 6}
// The `Init` function reorders the numbers into a heap
heap.Init(nums)
// The slice is now reordered to conform to the heap property
fmt.Println(nums)
}
Output:
&[1 1 1 5 3 4 2 6]
Note that the number is your are not necessarily in order. Rather they are arranged as a heap data structure which would look something like this:
Adding and Removing Elements From the Heap
We can use the heap.Pop function to get the current minimum (or maximum, for a max heap) element from a heap. This element is also removed from the heap, and the heap is restructured to maintain the heap property.
For example, if we want to remove the minimum element from nums:
min := heap.Pop(nums)
fmt.Println("min: ", min, " heap: ", *nums)
Output:
min: 1 heap: [1 3 1 5 6 4 2]
We can see that the heap is now restructured with the minimum value (1) removed:
To add an element to the heap we can use the heap.Push function along with the value of the element that we want to add:
heap.Push(nums, 5)
fmt.Println("heap: ", *nums)
Output:
heap: [1 3 1 5 6 4 2 5]
Note that both adding and removing an element take place in O(log(n)) time, so using a heap is a good way to maintain a priority queue of elements that can be added and removed efficiently.
Heap Sort
Since the pop function always returns the minimum element, we can continuously call the pop function for all elements in the heap to give us a sorted list of numbers:
sortedNums := []int{}
for nums.Len() > 0 {
sortedNums = append(sortedNums, heap.Pop(nums).(int))
}
fmt.Println("sorted values: ", sortedNums)
Output:
sorted values: [1 1 2 3 4 5 5 6]
The sort operation has to pop each item for all items, so the overall time complexity is O(n*log(n)).
Creating a Heap with Custom Struct Types
So far, we used an int heap, but we can create heaps from lists of any type, as long as we know the ordering property we want.
For example, consider a list of structs with a string field and a date field:
// Each holiday has a name and a date
type Holiday struct {
name string
date time.Time
}
// for convenience, we can create a string representation to print the holiday instance
func (h Holiday) String() string {
return "[" + h.name + ", " + h.date.Format("Jan 2") + "]"
}
// We can create a new type to represent a list of holidays
type Holidays []Holiday
Now there are two ways we can sort a list of Holiday instances:
- By their name, in alphabetical order
- By the calendar date
First, let’s implement the heap interface for the Holidays type to order by calendar date:
// Most of the changes here would be in the `Less` method
func (h Holidays) Less(i, j int) bool {
// Here, we choose to order by comparing the calendar dates of
// two holiday instances
return h[i].date.Before(h[j].date)
}
// The remaining methods are more or less the same as the previous examples
func (h Holidays) Len() int { return len(h) }
func (h Holidays) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *Holidays) Push(x interface{}) {
*h = append(*h, x.(Holiday))
}
func (h *Holidays) Pop() interface{} {
old := *h
n := len(old)
x := old[n-1]
*h = old[0 : n-1]
return x
}
Now we can use the heap methods to create our Holidays instance:
holidays := &Holidays{}
heap.Init(holidays)
heap.Push(holidays, Holiday{name: "Christmas", date: time.Date(2023, time.December, 25, 0, 0, 0, 0, time.Local)})
heap.Push(holidays, Holiday{name: "Labour Day", date: time.Date(2023, time.May, 1, 0, 0, 0, 0, time.Local)})
heap.Push(holidays, Holiday{name: "Diwali", date: time.Date(2023, time.November, 23, 0, 0, 0, 0, time.Local)})
fmt.Println("holidays: ", holidays)
Output:
holidays: &[[Labour Day, May 1] [Christmas, Dec 25] [Diwali, Nov 23]]
If we want to now order our holidays alphabetically, we need to change the Less method to use the name field for comparison:
// this will return true if holiday i is before holiday j in the dictionary
func (h Holidays) Less(i, j int) bool {
return strings.Compare(h[i].name, h[j].name) < 0
}
If we run the same heap example, our new output will be:
holidays: &[[Christmas, Dec 25] [Labour Day, May 1] [Diwali, Nov 23]]
If you want to see the code for all the examples in this post, you can view it on Github or run the live example on the Go playground
