C++でブラックリストを除外したランダム選択を実装する方法

範囲 [0, N) に属する一意な整数を格納したブラックリスト B があるとします。ここで求めたいのは、[0, N) の範囲からブラックリストに含まれていない整数を一様な確率で返す関数です。さらに、rand() の呼び出し回数をできるだけ減らすことで、この関数を最適化することも目標とします。例えば、入力として次のようなデータが与えられたケースを考えてみましょう。

解法のアプローチ

この問題を効率よく解く鍵となるのが「マッピング(対応表)」の考え方です。まず、実際に選択対象となる有効な数値の個数 M を事前に計算しておきます。そして、pick() が呼ばれた際には [0, M) の範囲でのみ乱数を生成します。もし生成された乱数がブラックリスト内の数値と一致した場合は、コンストラクタで事前に構築した対応表を使って、別の有効な数値へ置き換えて返します。

この仕組みにより、pick() を何度呼び出しても rand() はわずか1回で済むため、パフォーマンスの面で大きなメリットがあります。

アルゴリズムの手順

具体的な処理の流れは以下の通りです。

1. マップ(連想配列)を1つ用意します。
2. N と配列 v で初期化を行います。
3. i := 0 から v のサイズ未満の間、i を1ずつ増やしながら次を繰り返します。
   　・v[i] < N の場合:m[v[i]] := -1 とする
4. M := N − マップのサイズ とします。
5. n := v のサイズ とします。
6. 再び i := 0 から v のサイズ未満の間、i を1ずつ増やしながら次を繰り返します。
   　・v[i] < M の場合:
   　　1. N を1減らす
   　　2. N がマップ内に存在する間、N を減らし続ける
   　　3. m[v[i]] := N と登録する
7. pick() 関数を定義します。
8. x := 乱数 mod M を計算します。
9. x がマップに存在すれば m[x] を、存在しなければそのまま x を返します。

C++による実装例

それでは、実際のコードを見て理解を深めましょう。

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
    int M;
    map <int,int> m;
    Solution(int N, vector<int>& v) {
        for(int i = 0; i < v.size(); i++){
            if(v[i] < N) m[v[i]] = -1;
        }
        M = N - (int)(m.size());
        int n = v.size();
        for(int i = 0; i < v.size(); i++){
            if(v[i] < M){
                while(m.count(--N));
                m[v[i]] = N;
            }
        }
    }
    int pick() {
        int x = rand() % M;
        return m.count(x)? m[x] : x;
    }
};
main(){
    vector<int> v = {2};
    Solution ob(4,v);
    cout << (ob.pick()) << endl;
    cout << (ob.pick()) << endl;
    cout << (ob.pick()) << endl;
}

入力

N = 4、配列 [2]

出力

1
1
0

動作のポイント

この例では N = 4 であり、範囲 [0, 4) からブラックリスト {2} を除いた {0, 1, 3} のいずれかが返されます。コンストラクタでは、M 未満のブラックリスト要素に対して、範囲の後半側にある未使用の数値への対応付けを行っています。これにより、pick() 内で生成した乱数がブラックリスト要素と一致した場合でも、即座に有効な数値へ差し替えられる仕組みです。出力結果は乱数を使用しているため、実行ごとに変化する点にも注意してください。

計算量について

コンストラクタでは、ブラックリストのサイズを B とした場合に O(B log B) の時間計算量がかかります(マップ操作のため)。一方、pick() は O(log B) で動作し、乱数生成は毎回1回だけで済むため、頻繁に乱数を選択する場面で非常に効率的です。