Posted at

[Go言語] sort.Searchでランキングを作ってみた

More than 5 years have passed since last update.

3連休なので、sort.Search関数を使ってみました。この関数は、ソート済みのスライスや配列をバイナリサーチするのに使います。[0, n)の範囲を探索し、ftrueとなる最初のiを返します。f[0, n)の範囲で、f(i+1) == trueであればf(i+1) == trueであるような関数でなければなりません。

func Search(n int, f func(i int) bool) int

ちょっと分かり辛いので、具体例を見てみましょう。パッケージドキュメントにある例です。


example1.go

x := 23

i := sort.Search(len(data), func(i int) bool { return data[i] >= x })
if i < len(data) && data[i] == x {
// x is present at data[i]
} else {
// x is not present in data,
// but i is the index where it would be inserted.
}

検索対象はdataというスライスで、23を検索しています。

もし、sort.Searchに渡した関数は、data[i] >= xtrueとなる最初のiを返します。dataが昇順に並んでいるとすると、xすなわち23以上となる最初の添字を返すというわけです。もし、23dataにない場合は、len(data)+1が返ります。そのため、if文でやっているのは、i < len(data)fを満たす添字iがあった)かつdata[i]23であった、つまりdataの中に23があるかないかを判定してます。

このsort.Searchが便利なのは、添字で扱える点です。Go言語にはジェネリクスは無いですが、この方法であれば、添字だけを考えればいいので、どんな型のスライスや配列にも使えます。むしろ、スライスや配列である必要すらありません。math/randパッケージのrand.Permでも添字を使った方法が使われています。

せっかくなので、構造体を使った例をあげてみましょう。

http://play.golang.org/p/5u2H9WDgCt


ranking.go

package main

import (
"fmt"
"sort"
)

type Item struct {
id string
score int
}

func add(items []Item, item Item) []Item {
fmt.Println(items, "<=", item)

i := sort.Search(len(items), func(i int) bool {
return items[i].score < item.score ||
(items[i].score == item.score && items[i].id > item.id)
})

items = append(items, Item{})
copy(items[i+1:], items[i:])
items[i] = item

return items
}

func main() {
items := []Item{}
items = add(items, Item{"A", 100})
items = add(items, Item{"C", 200})
items = add(items, Item{"B", 200})
items = add(items, Item{"D", 50})
items = add(items, Item{"E", 100})
fmt.Println(items)
}


要素をスコア順(降順)に並べてみる例です。

同スコアだった場合はIDで辞書順になるようにしています。

return items[i].score < item.score || (items[i].score == item.score && items[i].id > item.id)の部分ですね。

このプログラムでは、同じIDのitemが追加された場合を考慮していません。考慮したければ、一度削除する必要があります。

http://play.golang.org/p/SVBZPmf8Lq


example2.go

package main

import (
"fmt"
"sort"
)

type Item struct {
id string
score int
}

func (item *Item) String() string {
return fmt.Sprintf("{%s:%d}", item.id, item.score)
}

type Ranking struct {
items []*Item
m map[string]*Item
}

func (r *Ranking) String() string {
return fmt.Sprintf("%v", r.items)
}

func (r *Ranking) Add(item *Item) {
fmt.Println(r, "<=", item)

// remove before insert
if old, ok := r.m[item.id]; ok {
i := sort.Search(len(r.items), func(i int) bool {
return r.items[i].score < old.score ||
(r.items[i].score == old.score && r.items[i].id >= old.id)
})
copy(r.items[i:], r.items[i+1:])
r.items[len(r.items)-1] = nil
r.items = r.items[:len(r.items)-1]
delete(r.m, old.id)
}

i := sort.Search(len(r.items), func(i int) bool {
return r.items[i].score < item.score ||
(r.items[i].score == item.score && r.items[i].id > item.id)
})

r.items = append(r.items, nil)
copy(r.items[i+1:], r.items[i:])
r.items[i] = item
r.m[item.id] = item
}

func main() {
ranking := &Ranking{make([]*Item, 0, 0), make(map[string]*Item)}
ranking.Add(&Item{"A", 100})
ranking.Add(&Item{"C", 200})
ranking.Add(&Item{"B", 200})
ranking.Add(&Item{"D", 50})
ranking.Add(&Item{"E", 100})
ranking.Add(&Item{"A", 200})
fmt.Println(ranking)
}


ちょっと頑張ってみました。

ランキングを構造体にしてみました。

古いitemmapから取得して、そのスコアを元に一度検索して、削除しています。

わりと簡単にsort.Searchが使えました。面白いですね。

ちなみに、スライスの操作はSliceTricksを参考にしました。