goでheapを実装する

heapとは優先度付きキューの一種。親node > 子node又は親node < 子nodeな制約を持った木構造です。

https://upload.wikimedia.org/wikipedia/commons/6/60/Binary_heap_indexing.png

出典 wikipedia

画像で見ると分かり易いですね。配列の中で一番小さい値がroot node。そこから深くなる毎に値が大きくなっています。左右の子ノード同士に大小関係はありません。

似たデータ構造に二分探索木がありますが、あちらは「左の子孫ノードの値 < 親ノードの値 < 右の子孫ノードの値」な制約をもった木構造です。

https://upload.wikimedia.org/wikipedia/commons/thumb/d/da/Binary_search_tree.svg/600px-Binary_search_tree.svg.png

出典 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 Pop5つのメソッドを定義したHeap構造体を作成し、初期化メソッドheap.Initの引数として渡す事でroot nodeを最小としたheap木を作成してくれます。

heap.Push,heap.Popは一見配列の要素を1つ増減させるだけのメソッドに見えますが、内部では配列内を走査し、親node < 子nodeの制約の元indexを並べ変えています。

標準ライブラリを使った実装パターンが分かった所で少し深堀りしていきましょう。

深堀

この動画を見てください!

www.youtube.com

FIN

許してください。手抜きではないです。どんなに文字を尽してもこの動画よりわかりやすい解説はできないんです、、、

アルゴリズムの様に動的にデータが移動する概念に関しては動画を見た方が絶対理解が捗ると思います。この方の動画全部面白いので他のも見て行ってください、、、

この後は蛇足ではありますが動画内の実装内容と標準ライブラリのインターフェースが若干解離している為、文面での解説を加えながら少し解説します

Push と Popの内部処理

先ほども言及しましたがheap.Pushheap.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で都度ソートしなおします

具体的なメソッドの挙動としては

  1. 頂点の値(root node)を最後尾の値にすげ替える
  2. 最後尾の値を削除
  3. 新しく代入されたroot nodeと二つの子nodeの内小さい方の値を比較。root nodeの方が大きい場合swapを実行
  4. これを繰り返す

こうして親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パッケージを使わない分実行速度の面でも多少優れています)