C++で重み付きランダム選択を実装する方法

正の整数の配列 w が与えられ、w[i] がインデックス i の重みを表しているとします。このとき、重みに比例した確率でインデックスをランダムに選択する関数 pickIndex() を定義する必要があります。

例えば、入力が [1, 3] の場合、pickIndex() を5回呼び出すと、結果は 0, 1, 1, 1, 0 のようになります。インデックス 1 の重みが 3 であるため、インデックス 1 が選ばれる確率はインデックス 0 の3倍になります。

解決策のアプローチ

この問題を解くために、以下の手順に従います。

• 配列 v を定義します。
• コンストラクタで以下のように初期化します。
• n := w[0] とします。
• i を 1 から w のサイズまでループします。
  • w[i] := w[i] + w[i - 1]（累積和を計算）
  • n := w[i]
• v = w として累積和の配列を保存します。

pickIndex() は以下のように動作します。

• 乱数 r を生成し、r mod v の最後の要素 を計算します。
• v の中から r 以上の最小の要素を二分探索で見つけ、その位置を返します。

この手法のポイントは、累積和配列を使うことで、各インデックスが選ばれる確率がその重みに正確に比例することを保証できる点です。また、upper_bound による二分探索を用いることで、選択処理を O(log n) の時間計算量で効率的に行えます。

実装例

以下の実装を見て、より深く理解しましょう。

#include <bits/stdc++.h>
using namespace std;
class Solution {
    public:
    int n;
    vector <int> v;
    Solution(vector<int>& w) {
        srand(time(NULL));
        n = w[0];
        for(int i = 1; i < w.size(); i++){
            w[i] += w[i - 1];
            n = w[i];
        }
        v = w;
    }
    int pickIndex() {
        return upper_bound(v.begin(), v.end(), rand() % v.back()) - v.begin();
    }
};
main(){
    vector<int> v = {1,3};
    Solution ob(v);
    cout << (ob.pickIndex()) << endl;
    cout << (ob.pickIndex()) << endl;
    cout << (ob.pickIndex()) << endl;
    cout << (ob.pickIndex()) << endl;
    cout << (ob.pickIndex()) << endl;
}

入力

[1, 3] で初期化し、pickIndex を5回呼び出します。

出力

1
1
1
1
0

この出力例では、重み 3 を持つインデックス 1 が多く選ばれていることが確認できます。ただし、ランダムな選択であるため、実行するたびに結果は異なる可能性がある点に注意してください。

計算量のまとめ

• コンストラクタ： O(n) — 累積和の計算に配列全体を一度走査します。
• pickIndex()： O(log n) — 二分探索（upper_bound）による選択を行います。
• 空間計算量： O(n) — 累積和を格納する配列が必要です。