蓄水池抽样

蓄水池抽样算法

蓄水池抽样算法随机算法的一种,用来从 N 个样本中随机选择 K 个样本,其中 N 非常大(以至于 N 个样本不能同时放入内存)或者 N 是一个未知数。其时间复杂度为 O(N)

解法

算法首先创建一个长度为 k 的数组(蓄水池)用来存放结果,初始化为 S 的前 k 个元素。然后从 k+1 个元素开始迭代直到数组结束,在 S 的第 i + 1个元素,算法生成一个随机数 𝑗[0,𝑖], 如果 j < k, 那么蓄水池的第 j 个元素被替换为 S 的第 i 个元素。

vector<int> ReservoirSampling(vector<int> data, int k) {
    int n = data.size();
    assert(k <= n);
    // init: fill the first k elems into reservoir
    vector<int> reservoir(data.begin(), data.begin() + k);
    for (int i = k; i < data.size(); i++) {
        int j = rand() % (i + 1);// inclusive range [0, i]
        if (j < k) {
            reservoir[j] = data[i];
        }
    }
    return reservoir;
}

证明

证明:首先,对于任意的 i,第 i 个元素进入蓄水池的概率为 k / i;而在蓄水池内每个元素被替换的概率为 1 / k; 因此在第 i 轮第j个元素被替换的概率为 (k / i ) * (1 / k) = 1 / i。 接下来用数学归纳法来证明,当循环结束时每个元素进入蓄水池的概率为 k / n.

假设在 (i-1) 次迭代后,任意一个元素进入 蓄水池的概率为 k / (i-1)。有上面的结论,在第 i 次迭代时,该元素被替换的概率为 1 / i, 那么其不被替换的概率则为 1 – 1/i = (i-1)/i;在第i 此迭代后,该元素在蓄水池内的概率为 k / (i-1) * (i-1)/i = k / i. 归纳部分结束。

因此当循环结束时,每个元素进入蓄水池的概率为 k / n. 命题得证。

参考

蓄水池抽样算法 (Reservoir Sampling Algorithm)

0