ガチャシステムを二分探索で作る

概要

以下のような確率分布があるとして、これに従ってランダムに要素を1つ決定したい。

実装のソースコードは記事の最後に記載する。運用先としてUnity等が真っ先に思い付いたため、C#での実装とした。
この機能に名前があるのか知らないが、近年ソーシャルゲームで流行しているガチャに似ているため、ガチャシステムと呼ぶ。

考え方

入力として、$N$ 個の要素からなる確率分布の数列 $A_1,A_2,⋯,A_N$ を与える。

double[] odds = { 0.5, 0.3, 0.15, 0.05 };

まず、$[0, \sum_{i=1}^{N} A_i]$ の範囲でランダムな実数 $r$ を1つ決定する。

Random rnd = new Random();
double r = rnd.NextDouble();

次に、$\sum_{i=1}^{k} A_i \leq r < \sum_{i=1}^{k+1} A_i$ となる整数 $k$ を求める。
簡単に思いつく方法は線形探索だが、今回は高速化を試みて二分探索で実装してみる。　※二分探索の説明は割愛する。
なお、確率分布の部分和を計算しなければいけないため、実装全体としては $O(N)$ である。$O(\log N)$ ではない。

とりあえず関数のシグネチャを定義。
引数として、先ほどの確率分布 $A_1,A_2,⋯,A_N$ 、そしてランダムに決定した値 $r ; (0 \leq r \leq \sum_{i=1}^{N} A_i)$ を取る。
戻り値は $r$ に対応するインデックス $k$ にすればよい。

public int BinarySearchInOdds(double[] odds, double r) {
    // oddsの確率分布から、rの値に対応するインデックスを返す
}

この部分の実装は、一般的な二分探索とあまり変わらないため説明を割愛する。
考え方としては以上のような形になる。

実装

BinarySearchInOdds() メソッドが、先ほど省略した「確率分布から、rの値に対応するインデックスを返す機能」を実装されているものである。
Draw() メソッドに確率分布を渡すことでお手軽にガチャを引くことができる。
仕様が分かるように Main() を書いたので、実行結果も含めてぜひ見ていってほしい。

using System;
using System.Collections.Generic;
using System.Linq;

public class Gacha {

    // 確率分布の中で、 rに対応する要素のインデックスを返す
    public static int BinarySearchInOdds(double[] odds, double r) {
        if (r < 0) return -1;
        if (odds.Length == 0) return -1;

        // 確率分布の部分和をあらかじめ求めておく
        double[] sum_from_left = new double[odds.Length];
        double[] sum_from_right = new double[odds.Length];
        double odd_total = odds.Sum();
        sum_from_left[0] = 0;
        sum_from_right[0] = odd_total;
        for (int i = 1; i < odds.Length; i++) {
            sum_from_left[i] = sum_from_left[i-1] + odds[i-1];
            sum_from_right[i] = sum_from_right[i-1] - odds[i-1];
        }

        int left = 0;
        int right = odds.Length - 1;
        int mid = left + (right - left) / 2;
        double L_mid = sum_from_left[mid];

        while (right >= left) {
            double R_mid = L_mid + odds[mid];
            if ((L_mid <= r && r < R_mid) || (r == R_mid && mid == odds.Length - 1)) {
                // キーがmid番目の要素の範囲内にある場合
                return mid;
            } else if (r < L_mid) {
                // キーがmid番目の要素の範囲より小さい範囲にある場合
                right = mid - 1;
                mid = left + (right - left) / 2;
                R_mid = L_mid - (odd_total - sum_from_left[mid + 1] - sum_from_right[right + 1]);
                L_mid = R_mid - odds[mid];
            } else {
                // キーがmid番目の要素の範囲より大きい範囲にある場合
                left = mid + 1;
                mid = left + (right - left) / 2;
                if (mid >= odds.Length) break;
                L_mid = R_mid + (odd_total - sum_from_left[left] - sum_from_right[mid]);
                R_mid = L_mid + odds[mid];
            }
        }
        return -1;
    }

    // 確率分布にしたがって、ランダムに決定した要素のインデックスを返す
    public static int Draw(double[] odds) {
        Random rnd = new Random();
        double r = rnd.NextDouble() * odds.Sum();
        return BinarySearchInOdds(odds, r);
    }

    public static void Main() {
        double[] odds1 = { 0.5, 3.0, 0, 2.0, 1.5, 0.1, 0.9 };
        Console.WriteLine("odds1 = [" + string.Join(", ", odds1) + "]");
        Console.WriteLine("r = 0.0  --->  " + BinarySearchInOdds(odds1, 0.0));
        Console.WriteLine("r = 0.1  --->  " + BinarySearchInOdds(odds1, 0.1));
        // rがちょうど境目にある時、大きいほうのインデックスが返ってくる
        Console.WriteLine("r = 0.5  --->  " + BinarySearchInOdds(odds1, 0.5));
        Console.WriteLine("r = 2.0  --->  " + BinarySearchInOdds(odds1, 2.0));
        Console.WriteLine("r = 4.0  --->  " + BinarySearchInOdds(odds1, 4.0));
        Console.WriteLine("r = 6.0  --->  " + BinarySearchInOdds(odds1, 6.0));
        Console.WriteLine("r = 7.08  --->  " + BinarySearchInOdds(odds1, 7.08));
        Console.WriteLine("r = 7.7  --->  " + BinarySearchInOdds(odds1, 7.7));
        // rが確率分布の合計と等しい場合、最後の要素のインデックスが返ってくる
        Console.WriteLine("r = 8.0  --->  " + BinarySearchInOdds(odds1, 8.0));
        // rが範囲外の場合、-1が返ってくる
        Console.WriteLine("r = 10.0  --->  " + BinarySearchInOdds(odds1, 10.0));
        // rが負の場合も同様
        Console.WriteLine("r = -1.0  --->  " + BinarySearchInOdds(odds1, -1.0));
        Console.WriteLine("");

        double[] odds2 = { 1.0 };
        Console.WriteLine("odds2 = [" + string.Join(", ", odds2) + "]");
        // 確率分布の要素が1つでも問題ない
        Console.WriteLine("r = 0.0  --->  " + BinarySearchInOdds(odds2, 0.0));
        Console.WriteLine("r = 0.5  --->  " + BinarySearchInOdds(odds2, 0.5));
        Console.WriteLine("r = 1.0  --->  " + BinarySearchInOdds(odds2, 1.0));
        Console.WriteLine("");

        double[] odds3 = {};
        // 確率分布の要素がない場合、-1が返ってくる
        Console.WriteLine("odds3 = [" + string.Join(", ", odds3) + "]");
        Console.WriteLine("r = 0.0  --->  " + BinarySearchInOdds(odds3, 0.0));
        Console.WriteLine("");

        double[] odds4 = { 0.5, 0.35, 0.15, 0.05 };
        Console.WriteLine("odds4 = [" + string.Join(", ", odds4) + "]");
        // 100万回ガチャを引いて、排出確率を調べてみる
        int[] receivedCounts = new int[4];
        for (int i = 0; i < 1000000; i++) {
            int received = Draw(odds4);
            receivedCounts[received]++;
        }
        for (int i = 0; i < 4; i++) {
            Console.WriteLine("要素" + i + "  " + receivedCounts[i] + " / 1000000  ( " + receivedCounts[i] / 10000 + " %)");
        }
    }

}

実行結果

odds1 = [0.5, 3, 0, 2, 1.5, 0.1, 0.9]
r = 0.0  --->  0
r = 0.1  --->  0
r = 0.5  --->  1
r = 2.0  --->  1
r = 4.0  --->  3
r = 6.0  --->  4
r = 7.08  --->  5
r = 7.7  --->  6
r = 8.0  --->  6
r = 10.0  --->  -1
r = -1.0  --->  -1

odds2 = [1]
r = 0.0  --->  0
r = 0.5  --->  0
r = 1.0  --->  0

odds3 = []
r = 0.0  --->  -1

odds4 = [0.5, 0.35, 0.15, 0.05]
要素0  456796 / 1000000  ( 45 %)
要素1  333225 / 1000000  ( 33 %)
要素2  144646 / 1000000  ( 14 %)
要素3  65333 / 1000000  ( 6 %)

実行速度も計測した。

N = 10^2  ->  2 ms
N = 10^3  ->  2 ms
N = 10^4  ->  2 ms
N = 10^5  ->  4 ms
N = 10^6  ->  18 ms
N = 10^7  ->  173 ms
N = 10^8  ->  1683 ms

$N=10^5$ 程度までは処理時間のオーダーがほぼ変わらない模様。二分探索による高速化が多少活きた...かな？

これでいつでもガチャ引き放題。