Go standard library has container/heap package which can be used to implement min/max heap and priority queue using this interface. It was implemented long time ago and there is a proposal to use Generics.
C++ STLのContainersとAlgorithms - sambaiz-net
type Interface interface {
sort.Interface
Push(x any) // add x as element Len()
Pop() any // remove and return element Len() - 1.
}
Try to use IntHeap in the document. In Pop(), it takes the value pointed by the pointer, processes the slice, and assigns it to itself.
package main
import (
"container/heap"
"fmt"
)
type IntHeap []int
func (h IntHeap) Len() int { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h IntHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
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
}
func main() {
h := &IntHeap{2, 1, 5}
heap.Init(h)
heap.Push(h, 3)
fmt.Printf("minimum: %d\n", (*h)[0]) // 1
for h.Len() > 0 {
fmt.Printf("%d ", heap.Pop(h)) // 1 2 3 5
}
}
Let’s see how these functions are called in heap.go.
Init() calls down() function some times.
func Init(h Interface) {
// heapify
n := h.Len()
for i := n/2 - 1; i >= 0; i-- {
down(h, i, n)
}
}
down() takes Interface h, index i0, and number of elements n, and if h[i] is larger than the smaller of the left child h[2i+1] and the right child h[2i+2], it swaps them and repeat that until they are not swapped. In other words, it moves down until it reaches a state where it has a value smaller than any of its children.
func down(h Interface, 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 && h.Less(j2, j1) {
j = j2 // = 2*i + 2 // right child
}
if !h.Less(j, i) {
break
}
h.Swap(i, j)
i = j
}
return i > i0
}
down() is needed only for non-leaf nodes, so it is called n/2 times, and nodes closer to the leaves are many and moved fewer times than the nodes closer to the root, so initialize can be executed in O(n). This is equivalent to heapq.heapify() which converts Python list to min heap.
From the following example, h[1]=7 > min(3, 6) so it is swapped with 3, and no further swaps are needed, and then h[0]=2 is also already smaller than any of its children, so the binary heap is complete.
h := [2, 7, 5, 3, 6]
2
/ \
7 5
/ \
3 6
down(h, 1, 5) => [2, 3, 5, 7, 6]
down(h, 0, 5) => no change
2
/ \
3 5
/ \
7 6
Pop() reconstructs the tree by bringing the tail value to popped root position and downing it. At this time, the value at the end is duplicated, so it must be removed by Pop() of Interface.
func Pop(h Interface) any {
n := h.Len() - 1
h.Swap(0, n)
down(h, 0, n)
return h.Pop()
}
Push() moves up a value that is appended to the tail by up().
func Push(h Interface, x any) {
h.Push(x)
up(h, h.Len()-1)
}
func up(h Interface, j int) {
for {
i := (j - 1) / 2 // parent
if i == j || !h.Less(j, i) {
break
}
h.Swap(i, j)
j = i
}
}
Remove() brings the tail value to root position same as Pop(), but since the value at the tail is not necessarily the largest, it also tries to move it down() and up(). Fix() called when the value is updated also does the same processing.
func Remove(h Interface, i int) any {
n := h.Len() - 1
if n != i {
h.Swap(i, n)
if !down(h, i, n) {
up(h, i)
}
}
return h.Pop()
}
func Fix(h Interface, i int) {
if !down(h, i, h.Len()) {
up(h, i)
}
}