N要素の一様ランダムな整数集合
$ N要素の$ L以上$ R未満の整数からなる集合を一様ランダムに生成したいとき。
TL;DR
以下のコードを使うと、期待$ O(N)で生成できる!(noshi91氏によるアルゴリズム)
code:rust
use rand::prelude::*;
use std::collections::*;
fn gen_random_set(n: usize, l: i64, r: i64) -> Vec<i64> {
assert!(l < r && (r - l) as usize >= n);
let mut rng = thread_rng();
let mut generated =
if (r - l) as usize <= n * 2
let mut samples = (l .. r).collect::<Vec<_>>();
samples.shuffle(&mut rng);
samplesn.to_vec()
} else {
let mut generated = HashSet::new();
while generated.len() < n {
generated.insert(rng.gen_range(l .. r));
}
generated.into_iter().collect()
};
generated.sort();
generated
}
アルゴリズム
[0, N) から相異なる整数 M 個をサンプリングする場合
・M ≥ N/2 なら N 個全部並べて shuffle して前から M 個取る
・M < N/2 なら何度もサンプリングしながら hashset で今まで出たやつをスキップして、M 個揃うまでやる
とすると O(M) で可能。
以下、31536000氏によるアルゴリズム
概要
Fisher-Yatesのアルゴリズムを用いると、$ iステップ目では配列の先頭から$ i個目の要素が決まるため、$ Nステップ目まですればよい。
期待$ O(N)でできる。
code:rust
use std::collections::*;
use rand::prelude::*;
fn gen_random_set(n: usize, l: i64, r: i64) -> Vec<i64> {
assert!(l < r && (r - l) as usize >= n);
if (r - l) as usize == n {
return (l .. r).collect();
}
let mut rng = thread_rng();
// 昇順の Fisher-Yates シャッフルをする
// l .. r の長さ r - l の配列を仮想的に用意する
let mut a = HashMap::<i64, i64>::new();
let mut generated = Vec::with_capacity(n);
for i in 0 .. n as i64 {
let j = rng.gen_range(i .. r - l);
generated.push(a.get(&j).copied().unwrap_or(l + j));
let x = a.get(&i).copied().unwrap_or(l + i);
a.insert(j, x);
}
generated.sort();
generated
}
以下、最初に考案したアルゴリズムの供養
アルゴリズム
値の種類数を$ W = R - Lとおく。
生成した整数の集合を$ Sとする。
1. $ S \leftarrow \emptyとする。
2. $ t = 0, 1, \ldots, N - 1について、以下を行う:
2. 1. $ [0, W - t) に含まれる整数を一様ランダムに生成し、これを$ i とする。
2. 2. まだ$ S に含まれない$ L 以上の整数であって、小さい方から$ i 番目のものを$ x とする
2. 3. $ Sに$ xを追加する。
正当性
$ Sを集合ではなく数列とし、追加する場合は末尾に追加すると考える。
すると、各ステップでは$ W, W - 1, \ldots, W - (N - 1) 通りのうちから選ばれているため、$ S の種類数は$ \frac{W!}{(W - N)!} 通りとなる。
これは$ W個の整数から$ N個を選び並び替える場合の数と同じである。
最後に$ S をソートすると、各結果について$ N! 個の場合が同一視されるため、$ \frac{W!}{(W - N)! N!} = \binom{W}{N}通りとなる。
Rustによる実装
集合を表現するにあたり、当初std::collctions::BTreeSetを使おうとしたが、$ xより小さい要素の個数を高速に求めることができなかったため、im_rc::Vectorを使った。
im_rc::Vectorは平衡二分木によるリストであり、挿入や二分探索を$ O(\log N)で行うことができる。
このコードは最悪$ O(N \log N \log (R - L))で動作する。
code:rust
use rand::prelude::*;
use im_rc::*;
fn gen_random_set(n: usize, l: i64, r: i64) -> Vec<i64> {
let mut rng = thread_rng();
let mut a = Vector::new();
for t in 0 .. n as i64 {
// まだ a にない整数のうち、小さい方から i 番目 (0-index)
let i = rng.gen_range(0 .. r - l - t);
// これまで生成された整数で x より小さいものが i 個以上であるような最小の x を二分探索
let x = {
let mut ac = r - 1;
let mut wa = l - 1;
while (ac - wa).abs() > 1 {
let wj = (ac + wa) / 2;
let j = a.binary_search(&wj).unwrap_or_else(|x| x ) as i64;
// wj より小さいものは wj - l 個あり、そこからこれまでに生成された j 個を除く
if wj - l - j >= i {
ac = wj;
} else {
wa = wj;
}
}
ac
};
a.insert_ord(x);
}
a.into_iter().collect()
}
See Also