goでheapを実装する
heapとは優先度付きキューの一種。親node > 子node
又は親node < 子node
な制約を持った木構造です。
出典 wikipedia
画像で見ると分かり易いですね。配列の中で一番小さい値がroot node。そこから深くなる毎に値が大きくなっています。左右の子ノード同士に大小関係はありません。
似たデータ構造に二分探索木がありますが、あちらは「左の子孫ノードの値 < 親ノードの値 < 右の子孫ノードの値」な制約をもった木構造です。
出典 wikipedia
使い分けとしては
- root node以外の特定のnodeを検索したい場合に効果は二分探索木
- treeの中の最大 or 最小値を知りたい時はヒープ
標準ライブラリcontainer/heap
goは標準ライブラリcontainer/heapでheapの実装をサポートしています。
// IntHeap は,整数の最小ヒープです。 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 interface{}) { // Push と Pop はポインタレシーバを使っています。 // なぜなら,スライスの中身だけでなく,スライスの長さも変更するからです。 *h = append(*h, x.(int)) } func (h *IntHeap) Pop() interface{} { old := *h n := len(old) x := old[n-1] *h = old[0 : n-1] return x } func Example_intHeap() { h := &IntHeap{2, 1, 5} heap.Init(h) heap.Push(h, 3) fmt.Printf("minimum: %d\n", (*h)[0]) for h.Len() > 0 { fmt.Printf("%d ", heap.Pop(h)) } // Output: // minimum: 1 // 1 2 3 5 }
こちらはGoDoc
から拝借したサンプルコードです。
この様にLen Less Swap Push Pop
5つのメソッドを定義したHeap
構造体を作成し、初期化メソッドheap.Init
の引数として渡す事でroot nodeを最小としたheap木を作成してくれます。
heap.Push
,heap.Pop
は一見配列の要素を1つ増減させるだけのメソッドに見えますが、内部では配列内を走査し、親node < 子node
の制約の元indexを並べ変えています。
標準ライブラリを使った実装パターンが分かった所で少し深堀りしていきましょう。
深堀
この動画を見てください!
FIN
許してください。手抜きではないです。どんなに文字を尽してもこの動画よりわかりやすい解説はできないんです、、、
アルゴリズムの様に動的にデータが移動する概念に関しては動画を見た方が絶対理解が捗ると思います。この方の動画全部面白いので他のも見て行ってください、、、
この後は蛇足ではありますが動画内の実装内容と標準ライブラリのインターフェースが若干解離している為、文面での解説を加えながら少し解説します
Push と Popの内部処理
先ほども言及しましたがheap.Push
とheap.Pop
は単なる要素増減以外にも、配列内の整合性を保つ為に「並び替え」を行っています。それが動画内のmaxHeapifyUp
, maxHeapifyDown
に当たる処理です。
container/heap
ライブラリでは最小の値をrootとしている為メソッド名はminHeapifyUp
, minHeapifyDown
とした方が自然でしょうか
まずはPush
の挙動から見ていきましょう。
type Heap []int func (h *Heap) Push(key int) { // 値を追加 *h = append(*h, key) // 最後尾のindexを指定 h.minHeapifyUp(len(*h)-1) } func (h *Heap) minHeapifyUp(index int) { // 親nodeより小さければ for (*h)[parent(index)] > (*h)[index] { h.swap(parent(index), index) index = parent(index) } } func parent(i int) int { return (i-1)/2 } func (h *Heap) swap(i1, i2 int) { (*h)[i1], (*h)[i2] = (*h)[i2], (*h)[i1] }
サンプルのPush
メソッドに1行加えました。
minHeapifyUp
は新しく加えられた値が親の値より小さければその位置をswap
し続けるメソッドです。
例えば新しく加えられた値がその配列の最小要素だった場合、root nodeに達するまでswapし続けます。
parent
関数は親nodeのインデックスを算出する為の関数です。treeにおけるindex番号の親子関係は左の子node = 親node * 2 + 1
、右の子node = 親node * 2 + 2
という式で表せる為、親indexも親node = (子node-1) / 2
で割り出せます。
たとえばindex1の値の子要素を調べたい場合。
図が大雑把で申し訳ありませんが、丸の中が値、左がindex番号です。
この様に1の子要素二つのindexは12+1の3と12+2の4であることがわかると思います。
次にPop
の処理は少し複雑
func (h *Heap) Pop() int { if len(*h) == 0 { fmt.Println("cannot extract beacause array length is 0") return -1 } // 除外されるroot node extracted := (*h)[0] l := len(*h)-1 // 最後尾の値をroot nodeにすげ替える (*h)[0] = (*h)[l] // 最後尾の値を切り捨て (*h) = (*h)[:l] // root nodeと子nodeを比較。root node < 子nodeの小さい方ならswap h.minHeapifyDown(0) return extracted } func (h *Heap) minHeapifyDown(index int) { lastIndex := len(*h) -1 l, r := left(index), right(index) childToCompare := 0 // 左辺しか存在しない場合 for l <= lastIndex { if l == lastIndex { childToCompare = l } else if (*h)[l] < (*h)[r] { childToCompare = l } else { childToCompare =r } if (*h)[index] > (*h)[childToCompare] { h.swap(index, childToCompare) index = childToCompare l, r = left(index), right(index) } else { return } } } func left(i int) int { return i*2+1 } func right(i int) int { return i*2+2 }
Pop関数の挙動は大まかに2つ。配列内の最小値(root node)を排除して排除した値を返す事です。
ポイントはroot nodeを排除した後もheap木の制約を守り続ける点。minHeapifyDown
で都度ソートしなおします
具体的なメソッドの挙動としては
- 頂点の値(root node)を最後尾の値にすげ替える
- 最後尾の値を削除
- 新しく代入されたroot nodeと二つの子nodeの内小さい方の値を比較。root nodeの方が大きい場合swapを実行
- これを繰り返す
こうして親node < 子nodeの制約を守り続けます。
例題 D - Prefix K-th Max
最後に演習問題を一つ
(1,2,…,N) の順列 P=(P1,P2,…,PN)、および正整数 K が与えられます。 i=K,K+1,…,N について、以下を求めてください。 P の先頭 i 項のうち、K 番目に大きい値
簡単に要約すると
- 配列の中から要素数Kのスライスを切り出し、その中から最小の値を出力
- これを
K,K+1,…,N
まで繰り返す
解答
root nodeを最小値としたheap treeを構成します。
func main() { var n, k int var pp []int h := &Heap{} for _, v := range pp[:k-1] { h.Push(v) } now := 0 for i := k-1; i < n; i++ { cur := pp[i] if cur > now { h.Push(cur) now = h.Pop() } fmt.Println(, now) } } type Heap []int func (h *Heap) Push(key int) { *h = append(*h, key) h.minHeapifyUp(len(*h)-1) } func (h *Heap) Pop() int { if len(*h) == 0 { fmt.Println("cannot extract beacause array length is 0") return -1 } extracted := (*h)[0] l := len(*h)-1 (*h)[0] = (*h)[l] (*h) = (*h)[:l] h.minHeapifyDown(0) return extracted } func (h *Heap) minHeapifyUp(index int) { for (*h)[parent(index)] > (*h)[index] { h.swap(parent(index), index) index = parent(index) } } func (h *Heap) minHeapifyDown(index int) { lastIndex := len(*h) -1 l, r := left(index), right(index) childToCompare := 0 // 左辺しか存在しない場合 for l <= lastIndex { if l == lastIndex { childToCompare = l } else if (*h)[l] < (*h)[r] { childToCompare = l } else { childToCompare =r } if (*h)[index] > (*h)[childToCompare] { h.swap(index, childToCompare) index = childToCompare l, r = left(index), right(index) } else { return } } } func parent(i int) int { return (i-1)/2 } func left(i int) int { return i*2+1 } func right(i int) int { return i*2+2 } func (h *Heap) swap(i1, i2 int) { (*h)[i1], (*h)[i2] = (*h)[i2], (*h)[i1] }
見づらいので入出力のバッファリングは省略しています。
正直標準パッケージでも全然問題ありません!が、実際に実装してみるとやっぱり頭に入りますね(reflect
パッケージを使わない分実行速度の面でも多少優れています)