要素数$n$の配列からランダムに重複を許さずに$k$個を選びたい場合、例えば、配列をシャッフルしてから最初の$k$個を選ぶことで達成できます。
特に$k=1$のときはシャッフルをせずに、ランダムに$1$以上$n$以下の整数を生成し(配列の添え字が1始まりの場合)、そのインデックスの要素を返すことでも達成できます。
では、要素数が不明な場合や$n$がとても大きい場合はどうしたらいいでしょうか?
そのような場合に使えるアルゴリズムにReservoir samplingがあるので解説したいと思います。
参考文献
https://en.wikipedia.org/wiki/Reservoir_sampling
https://gregable.com/2007/10/reservoir-sampling.html
Reservoir sampling
擬似コードで書くと以下のようになります。
便宜上、入力の配列source
の要素数を$n$としましたが、source
の要素を逐一処理していくので$n$が不明でも実行できます。
入力: 要素数nの配列sourceと、n以下の整数k
出力: 要素数kの配列result
resultにsourceの最初のk個の要素をそれぞれ代入する
iをk+1からnまで増加させながら、以下を実行する (配列の添え字が1始まりとする)
rに1以上i以下のランダムな整数を代入する
rがk以下ならば、result[r]にsource[i]を代入する
配列をシャッフルするアルゴリズムFisher–Yates shuffleとよく似ているのがわかります。
Javaでの実装はこんな感じになります。長さが不明でも実装できることを明確にするためにイテレータを入力としました。
public <T> List<T> reservoirSampling(Iterator<T> source, int k) {
List<T> result = new ArrayList<>(k);
for(int i = 0; i < k && source.hasNext(); i++) {
result.add(source.next());
}
Random random = new Random();
int i = k+1;
while(source.hasNext()) {
T element = source.next();
int r = random.nextInt(i);
if(r < k) {
result.set(r, element);
}
i++;
}
return result;
}
具体例
source = [A, B, C, D]
、k = 3
を例にとります。期待する結果は、source
から重複しない3つの要素をそれぞれ等しい確率$\frac{3}{4}$で選ぶことです。
まず出力の配列がresult = [A, B, C]
で初期化されます。
次に乱数$1 \leqq r \leqq 4 $を生成します。要素D
が選ばれるのは$r \leqq k$のとき、つまり$r=1, 2, 3$のいずれかの場合で確率は$\frac{3}{4}$です。
$r=4$の場合は選ばれないことを意味し、確率は$\frac{1}{4}$です。
このとき要素A
がresult
に残る、つまりA
がD
によって上書きされない確率は以下の和事象の確率です。
-
D
が選ばれない確率$\frac{1}{4}$ -
D
が選ばれて、[A, B, C]
のうちB
またはC
を上書きする確率$ \frac{3}{4} \times \frac{2}{3}$
従ってA
がresult
に残る確率も$\frac{1}{4} + \frac{3}{4} \times \frac{2}{3} = \frac{3}{4} $であることがわかります。
要素B
C
についても同様なので、A
B
C
D
のいずれもが等しい確率$\frac{3}{4}$で選ばれることがわかります。
証明
$i \ge k$とし、入力配列の$i$番目の要素$x_i$の次の要素$x_{i+1}$を処理する直前の時点で、$i$番目以下の要素$ \{x_1, x_2, ..., x_i\} $のいずれもが等しい確率$ P_i = \frac{k}{i} $で選ばれていると仮定します。
$i=k$のとき$ P_k = 1 $であり、アルゴリズムの定義より、入力配列の最初の$k$個の要素が全て選ばれていることから成り立つことは明らかです。
$i=j$のとき$ P_j = \frac{k}{j} $が成り立つと仮定します。
$i=j+1$のとき$P_{j+1} = \frac{k}{j+1}$で成り立つことを示します。
まず要素$x_{j+1}$が選ばれる確率はアルゴリズムの定義より$\frac{k}{j+1}$です。
要素$x_{j+1}$を処理したあと$j$番目以下の要素$ x_l \in \{x_1, x_2, ..., x_j\} $が選ばれている確率は、
(A) 要素$x_{j+1}$を処理する前にすでに$x_l$がすでに選ばれていて、かつ
(B) $x_l$が$x_{j+1}$によって上書きされない確率です。
(A)の確率は、仮定より確率$ P_j = \frac{k}{j} $です。
(B)の確率は、$x_{j+1}$が確率$\frac{k}{j+1}$で選ばれ、かつ$k$個の中から$x_l$が選ばれる事象の、余事象の確率$1 - (\frac{k}{j+1} \frac{1}{k}) = 1 - \frac{1}{j+1} = \frac{j}{j+1}$です。
よって要素$x_l$が選ばれる確率は(A)と(B)の積事象の確率で、これも$\frac{k}{j}\frac{j}{j+1} = \frac{k}{j+1}$ であることがわかりました。
以上より、要素$x_{j+1}$と$x_l$がともに確率$\frac{k}{j+1}$で選ばれることから$P_{j+1} = \frac{k}{j+1}$であることがわかりました。
数学的帰納法より正しいことが示されました。