Go
golang

[Go]重み付き乱択アルゴリズムを整数だけで完結させる

ある要素集合から要素を重みを付けてランダムに選出することを「重み付き乱択アルゴリズム」というらしいのですが、今回はそれを浮動小数点数を使わずに整数のみで実装しようと思います。
ソースコードは本記事の末尾にまとめて記載しています。

まずは重みの定義ですが、1以上の正整数からなる配列で定義しておきます。

waits := []int{1, 2, 3, 4, 5, 6, 7, 8, 9}

もし0や負整数が含まれていた場合はうまく動作しないので必ず正整数だけで定義しておきましょう。
また、配列が空(=重みが1つも定義されていない)の場合も同様にうまく動作しません。
実際に使う場合は独自のエラーハンドリングを記載しておいた方が良いですね。

重みを定義した後は、どこからどこまでが「$i$番目の重みが属するセクション」かを決めるために境界値を定義していきます。
これはただ単に重みの累積和をとるだけでいいですね。
"セクション"と言うからには最初の境界値が0である点に注意が必要です。

boundaries := make([]int, len(waits)+1)
for i := 1; i < len(boundaries); i++ {
    boundaries[i] = boundaries[i-1] + waits[i-1]
}

そしてここからがメインの重み付き乱択アルゴリズムの説明になります。

乱数の生成はmath/randパッケージのrand.Intnを使います。
これは引数にint型の正整数nを与えることで[0,n-1]の正整数を返す関数です。
ただし、のちほどの重み位置の探索との兼ね合いで乱数を+1しておきます。
つまり乱数の生成区間を[1,n]にします。

x := rand.Intn(boundaryLast)+1

次に、生成した乱数が境界値で区切ったセクションの何番目か、つまり何番目の重みかを求めていきます。
境界値が単調増加列になっているため二分探索を用いて求められます。
「二分探索のアルゴリズムとか実装するのめんど!」と思うかもしれませんが、わざわざそれを実装する必要はなくて、sortパッケージのsort.SearchIntsを用いればいいです。

idx := sort.SearchInts(boundaries, x)

最後に、$i$番目の境界値を$b_i$として、生成した乱数を二分探索した結果その乱数が区間$[b_{i},b_{i+1}]$に含まれる場合、セクション位置としては$i$番目ですが二分探索の結果は$i+1$が返るので、その結果に-1しておきます。

最終的なソースコードは以下の通りです。

package main

import (
    "fmt"
    "math/rand"
    "sort"
    "time"
)

const N = 1000000

func main() {
    rand.Seed(time.Now().UnixNano())

    waits := []int{1, 2, 3, 4, 5, 6, 7, 8, 9}

    boundaries := make([]int, len(waits)+1)
    for i := 1; i < len(boundaries); i++ {
        boundaries[i] = boundaries[i-1] + waits[i-1]
    }

    boundaryLast := boundaries[len(boundaries)-1]
    counter := make([]int, len(waits))
    for i := 0; i < N; i++ {
        x := rand.Intn(boundaryLast)+1
        idx := sort.SearchInts(boundaries, x) - 1
        counter[idx]++
    }

    fmt.Println("|wait|expected|actual|")
    fmt.Println("|:-:|:-:|:-:|")
    for i := 0; i < len(waits); i++ {
        fmt.Printf("|%d|%f|%f|\n", waits[i], float64(waits[i]) / float64(boundaryLast), float64(counter[i]) / float64(N))
    }
}

100万回選出した結果、重みごとの選出率の期待値(expected)と実際値(actual)はかね近いことが分かります。
実行するごとにactualは異なります。

wait expected actual
1 0.022222 0.022258
2 0.044444 0.044653
3 0.066667 0.066533
4 0.088889 0.088676
5 0.111111 0.110854
6 0.133333 0.132967
7 0.155556 0.156268
8 0.177778 0.177938
9 0.200000 0.199853

上記ソースコードの注意点として、重みの大きさによっては境界値がintで収まらない場合があるので、オーバーフローに気をつける必要があります。その時はuint64など別の型で対処しましょう。