蓄水池抽样算法

算法简介

蓄水池抽样算法是随机算法的一种，用来从 N 个样本中随机选择 m 个样本，其中 N 非常大（以至于 N 个样本不能同时放入内存），且 N 是一个未知数。请问如何在只遍历一遍数据（O(N)）的情况下，能够随机选取出m个不重复的数据。

综上，可以提炼出以下三点：

1. 数据流长度N很大且不可知，所以不能一次性存入内存；
2. 时间复杂度为O(N)；
3. 随机选取m个数，每个数被选中的概率为m/N；

第1点限制了不能直接取N内的m个随机数，然后按索引取出数据。第2点限制了不能先遍历一遍，然后分块存储数据，再随机选取。第3点是数据选取绝对随机的保证。

基本原理

算法首先创建一个长度为 m 的数组（蓄水池）用来存放结果，初始化为 S 的前 m 个元素。然后从 m+1 个元素开始迭代直到数组结束（N)，在 S 的第 i 个元素，算法生成一个随机数 j∈[1,i]， 如果 j <= k， 那么蓄水池的第 j 个元素被替换为 S 的第 i 个元素。对应伪代码如下：

array R[m];    // result
integer i, j;

// fill the reservoir array
for each i in 1 to m do
    R[i] := S[i]
done;

// replace elements with gradually decreasing probability
for each i in m+1 to length(S) do
    j := random(1, i);   // important: inclusive range
    if j <= m then
        R[j] := S[i]
    fi
done

算法正确性证明

待证明结论：该算法保证每个元素以 m / n 的概率被选入蓄水池数组。

证明：
首先，对于任意的 i，第 i 个元素进入蓄水池的概率为 m / i；而在蓄水池内每个元素被替换的概率为 1 / m; 因此在第 i 轮第j个元素被替换的概率为 (m / i ) * (1 / m) = 1 / i。 接下来用数学归纳法来证明，当循环结束时每个元素进入蓄水池的概率为 m / n.
假设在 (i-1) 次迭代后，任意一个元素进入 蓄水池的概率为 m / (i-1)。有上面的结论，在第 i 次迭代时，该元素被替换的概率为 1 / i， 那么其不被替换的概率则为 1 - 1/i = (i-1)/i；在第i 此迭代后，该元素在蓄水池内的概率为 m / (i-1) * (i-1)/i = m / i. 归纳部分结束。

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

算法的实现

#include <bits/stdc++.h>
using namespace std;

// 蓄水池抽样算法 
vector<int> ReservoirSampling(vector<int> v, int n, int k) {
    vector<int> reservoirArray(k,0);
    // 先将蓄水池填满
	for (int i = 0; i < k; i++) {
		reservoirArray[i] = i;
	}
	
	// 从k+1开始选，判断要不要放到蓄水池中 
    for (int i = k; i < n; ++i) {
        int j = rand() % (i + 1); 
        //  j∈[1,i]，替换 reservoirArray_j 
        if (j < k) {
            reservoirArray[j] = v[i];
        }
    }
    return reservoirArray;
}

int main()
{
	// 先构建测试用数组[0,1,2,3,4,5,6,7,8,9] 
    vector<int> v = {0,1,2,3,4,5,6,7,8,9};
    
    srand((unsigned int)time(NULL));
    const int RUN_COUNT = 10000;
    int cnt[11] = {0};
    
	// 运行10000次，统计每个数字出现的次数 
    for (int i = 1; i <= RUN_COUNT; ++i) {
        vector<int> samples = ReservoirSampling(v, 10, 5);
        //  统计每次，每个数字出现的次数 
        for (int j = 0; j < samples.size(); ++j) {
            cnt[samples[j]]++;
        }
    }
	cout << "last" << endl;
	// 最后，看下每个数字选到的次数 
    for (int num = 0; num < 10; num++) {
        cout << num << " : \t" << cnt[num] << endl;
    }
    return 0;
}

运行结果如下

算法的局限性

蓄水池算法的基本假设是总的样本数很多，不能放入内存，暗示了选择的样本数 m 是一个与 n 无关的常数。然而在实际的应用中，m 常常与 n 相关，比如我们想要随机选择1/3 的样本 (m = n / 3)，这时候就需要别的算法或者分布式的算法。