C++で重複あり配列からランダムにインデックスを選択する方法（リザバーサンプリング解説）

問題概要

重複する要素を含む整数型配列が与えられたとき、指定したターゲット値が出現するインデックスの中から1つをランダムに選択して返すことを考えます。ターゲット値は必ず配列内に存在すると仮定できます。

例えば、配列が [1, 2, 3, 3, 3] の場合、pick(3) を呼び出すと、インデックス 2、3、4 のいずれかが等確率で返されます。

アルゴリズムの考え方

この問題は「リザバーサンプリング（Reservoir Sampling）」と呼ばれる手法を使うことで、配列を1回走査するだけで各候補を等確率で選択できます。手順は以下の通りです。

• 結果を格納する変数 ret を -1、カウンタ cnt を 1 で初期化します。
• i を 0 から配列 v のサイズまでループさせます。
  • v[i] が target と一致した場合：
    • 乱数を cnt で割った余りが 0 であれば、ret = i とします。
    • cnt を 1 増やします。
• 最後に ret を返します。

C++による実装例

以下に実際のC++コードを示します。

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
    vector<int> v;
    Solution(vector<int>& nums) {
        srand(time(NULL));
        v = nums;
    }
    int pick(int target) {
        int ret = -1;
        int cnt = 1;
        for(int i = 0; i < v.size(); i++){
            if(v[i] == target){
                if(rand() % cnt++ == 0) ret = i;
            }
        }
        return ret;
    }
};
main(){
    vector<int> v = {1,2,3,3,3};
    Solution ob(v);
    cout << (ob.pick(3));
}

入力

[1,2,3,3,3] で初期化
pick(3) を呼び出してランダムなインデックスを取得

出力

4
3
4
2

解説：なぜ等確率になるのか？

このアルゴリズムでは、target が k 回目に出現したとき、その時点でのカウンタ cnt の値は k になっています。rand() % k == 0 となる確率は 1/k なので、k 番目の要素が新しく選ばれる確率は 1/k、それ以前に選ばれていた要素が保持される確率は 1 − 1/k となります。これを帰納的に計算すると、target の出現回数を n としたとき、すべての出現位置が等確率 1/n で選択されることが証明できます。

この手法の利点は、事前に配列全体をスキャンしてインデックス一覧を作成しておく必要がなく、O(1) の追加メモリと O(n) の計算時間で処理できる点です。さらに、データがストリーム形式で逐次与えられる場合にも適用できるため、大量データやオンライン処理の場面で実務的に広く活用されています。