Goのcontainer/heap packageを読む
golangGoの標準パッケージに container/heap というものがあり、この interface を用いて min/max heap や priority queue を実装できる。昔からある package で、Generics を用いた proposal も上がっている。
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.
}
まずはドキュメントにある IntHeap を動かしてみる。Pop() では ポインタが指す値を取って slice を操作し自身に代入している。
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
}
}
これらの関数がどのように呼ばれているのかを見ていく。
まず、Init() すると down() という関数が n/2 回呼ばれる。
func Init(h Interface) {
// heapify
n := h.Len()
for i := n/2 - 1; i >= 0; i-- {
down(h, i, n)
}
}
この down() は上の Interface h と index i0, 要素数 nを受け取り、h[i] が左の子である h[2i + 1] と 右の子である h[2i + 2] の小さい方より大きい場合 Swapするというのを Swap しなくなるまで繰り返す。要は、いずれの子よりも小さな値を持つ状態になるまで下に動く。
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
}
次の例で考えると、まず h[1]=7 > min(3, 6) なので 3 と Swap され、それ以上下には動かないので続いて h[0]=2 を見るが、これは既にいずれの子よりも小さいのでそのまま二分ヒープが完成する。
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() では削除する root の代わりに末尾を持ってきて down することで木を再構成している。このとき配列の末尾の値が重複するので Interface の Pop() で取り除かなくてはならない。
func Pop(h Interface) any {
n := h.Len() - 1
h.Swap(0, n)
down(h, 0, n)
return h.Pop()
}
Push() では末尾に append された値を 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() は Pop() と同じく削除する要素を末尾で置き換えるが、ヒープの性質上末尾の値が最も大きいとは限らないので、もし down() で動かない場合は up() での移動も試みる。値が更新されたときの Fix() も同様の処理を行う。
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)
}
}