はじめまして、MF KESSAIでバックエンドエンジニアを担当しているshoheiです。
今回はGoの開発業務で触ったcontainer/heapパッケージについての話です。
container/heapとは
Overviewを読めば分かる通り、「ヒープ」データ構造とその操作関数を提供してくれるパッケージです。
使い方も非常に簡単でインタフェースを実装したスライスにheap.Pushしていくと、heap.Popでスライス内での最小値が返却されます。 スライスを二分木構造に展開してPop時にルートにあたる最小値を返しているんですね。よくあるユースケースとして以下に時刻のソートを実装してみます。
package main
import (
"container/heap"
"fmt"
"time"
)
type TimeHeap []time.Time
// Len スライス要素数
func (t TimeHeap) Len() int {return len(t)}
// Less 要素の比較式
func (t TimeHeap) Less(i, j int) bool {return t[i].Before(t[j])}
// Swap 要素の交換
func (t TimeHeap) Swap(i, j int) {t[i], t[j] = t[j], t[i]}
// Push 要素の追加
func (t *TimeHeap) Push(x interface{}) {*t = append(*t, x.(time.Time))}
// Pop 要素の抽出
func (t *TimeHeap) Pop() interface{} {
old := *t
n := len(old)
x := old[n-1]
*t = old[0 : n-1]
return x
}
func main() {
today := time.Now()
tomorrow := time.Now().AddDate(0, 0, 1)
yesterDay := time.Now().AddDate(0, 0,-1)
h := &TimeHeap{
today,
}
heap.Init(h)
heap.Push(h, tomorrow)
heap.Push(h, yesterDay)
for h.Len() > 0 {
fmt.Printf("%v\n", heap.Pop(h))
}
}- 処理結果
2020-02-08 16:42:20.828032 +0900 JST
2020-02-09 16:42:20.827907 +0900 JST m=+0.000075967
2020-02-10 16:42:20.827907 +0900 JST
上記は簡単な例ですがインタフェースを満たす実装をする事で任意の構造体スライスに対しても同様の機構を提供できます。
sort or heap
ここで疑問に感じるのは sort パッケージの存在、Goで一般的なソート処理といえば真っ先に思いつきます。単にスライスをソートしたいだけなら必要充分でしょう。
さっきの例をsortで書き直してみます。
package main
import (
"fmt"
"sort"
"time"
)
func main() {
today := time.Now()
tomorrow := time.Now().AddDate(0, 0, 1)
yesterDay := time.Now().AddDate(0, 0,-1)
dates := []time.Time{
today,
tomorrow,
yesterDay,
}
sort.Slice(dates, func(i, j int) bool {
return dates[i].Before(dates[j])
})
for _, date := range dates {
fmt.Printf("%v\n", date)
}
}heapを使った例よりも短く簡潔に書けますね。
では、heapの使い所はどのようなタイミングなのでしょうか。
container/heapの使い所
先程の例では単純に時刻をソートする例だけを考えていました。 すこしプログラムの要件を変えてみましょう。用意したスライスに要素が追加される事を想定します。
要素追加（Push）される都度、最小値を取得し要素から削除（Pop）される 事を要件としてみましょう。
まず、sortを利用した場合。
today := time.Now()
tomorrow := time.Now().AddDate(0, 0, 1)
yesterDay := time.Now().AddDate(0, 0, -1)
dates := []time.Time{
today,
tomorrow,
yesterDay,
}
rand.Seed(today.UnixNano())
for i := 0; i < 5; i++ {
addDate := today.AddDate(0, 0, rand.Intn(10))
// Push Sliceに要素追加
dates = append(dates, addDate)
// Pop 最小値を取得してSlice要素から除く
sort.Slice(dates, func(i, j int) bool {
return dates[i].Before(dates[j])
})
minDate := dates[0]
dates = dates[1:]
fmt.Printf("minDate: %s, addDate: %s, dates: %v\n",
minDate.Format("2006-01-02"),
addDate.Format("2006-01-02"),
dates,
)
}うーん、あまりいい実装には見えませんね。特に望ましくない点が要素追加の都度sortを実行している点です。 Goのsortはスライスの長さをnとした場合、計算量はO(n*log(n))となります。対して線形探索した場合の計算量はO(n)、sortを使わず愚直に線形探索したほうがパフォーマンス良さそうです。
次に、heapを利用した場合。
today := time.Now()
tomorrow := time.Now().AddDate(0, 0, 1)
yesterDay := time.Now().AddDate(0, 0,-1)
h := &TimeHeap{
today,
tomorrow,
yesterDay,
}
heap.Init(h)
rand.Seed(today.UnixNano())
for i := 0; i < 5; i++ {
addDate := today.AddDate(0, 0, rand.Intn(10))
// Push 要素追加
heap.Push(h, addDate)
// Pop 最小値を取得して要素から除く
minDate := heap.Pop(h)
fmt.Printf("minDate: %s, addDate: %s, dates: %v\n",
minDate.(time.Time).Format("2006-01-02"),
addDate.Format("2006-01-02"),
h,
)
}先程の例よりも簡潔に実装ができています。
heapの計算量はheap.InitがO(n)、heap.Pushとheap.PopがO(log(n)) となっています。 要素追加の都度発生する計算量はO(log(n))のため、処理効率の観点でもheapの方がパフォーマンスを期待できますね。
まとめ
sortに比べて少し日陰がちなcontainer/heapに焦点をあてユースケースを考えてみました。
今回は計算量での優位性に着目してsortとの比較をしていましたが、heapは他にも有効に利用できるケースがあります。代表的な利用例は 優先度付きキュー の実装でしょう。 ドキュメントのExampleでは優先度付きキューの実装例が提示されています。 優先度付きキューの実装が必要になった際はそちらも参照してみてください。
日々の開発においてソーティングの絡む実装が必要になるケースは多々発生します。 一辺倒にsortを使うのではなく、要件に応じたデータ構造やアルゴリズムを選択するようにしましょう。 今回はその選択肢の1つとしてcontainer/heapパッケージを紹介しました。
もっと詳しい話を聞きたい方、MF KESSAIのエンジニアの仕事に興味がある方、（処理ディティールに拘りたい方）、ぜひお気軽にMF KESSAIへ遊びにきてください！