蓄水池抽样算法

问题描述

问题：有一段很长整数流，不知道其长度，内存可能也不够存下所有数据。希望从这个整数流中抽样出$k$个整数，要求每个整数被抽样的概率相同。

这种问题称之为蓄水池抽样问题。

最大堆解决

如果一个数组是随机排列的，我们取前$k$个元素也是随机的，那么这$k$元素可以称之为一个随机抽样。

——非名家之言

依以上所言，我们只需要得到整数流的一个随机排列就可以了。但是我们连长度$n$都不知道，怎么随机排列？

我们一定要排列好所有元素吗？我们只需要抽样$k$个，因此我们只需要知道哪$k$个元素排在前面就可以。这是不是让你想起了这样一个问题：一堆数中取前$k$个最大的数。这个问题可以用堆解决，同样蓄水池抽样问题也可以用堆解决。

算法描述：维护一个大小为$k$的大顶堆，每次来一个数，随机生成一个key，如果当前元素的key比堆顶的元素要小，弹出堆顶元素，将当前元素和key加入堆中。最终堆里的元素是一个随机排列（因为key是随机生成的）里最小的前$k$个数。这$k$个数就是一个随机抽样。（也可以用小顶堆，那就是取最大的的$k$个数）

代码：

vector<int> reservoirSamplingByHeap(vector<int> &nums, int k) {
    priority_queue< pair<int, int> > heap;
    for (int i = 0; i < nums.size(); i++) {
        int r = rand(); // 随机生成一个key
        if (heap.size() < k) heap.push(make_pair(nums[i], r));
        else {
            if (heap.top().second > r) {
                heap.pop();
                heap.push(make_pair(nums[i], r));
            }
        }
    }
    vector<int> sample;
    while (!heap.empty()) {
        sample.push_back(heap.top().first);
        heap.pop();
    }
    return sample;
}

时间复杂度：$nlgk$

更快的解法

$nlgk$已经足够快，但是也有更快的解法，也相当巧妙。

算法描述：维护一个大小为$k$的蓄水池，初始将前$k$个元素放入蓄水池中，从第$k+1$个元素开始，每个元素以$\frac{k}{i}$的概率选中，随机替换蓄水池中的一个元素。通过这种方式每个元素被选取到的概率相等。

证明：

我们先从最简单的问题入手：$k=1$

则第$i$个元素被选取的概率是$\frac{1}{i}$，但是现在被选取了并不代表后来不会被替换掉，所以元素$i$最终被选取为被选取的概率乘上不被替换的概率：

$P(i)=\frac{1}{i}\times\frac{i}{i+1}\times\cdots\times\frac{n-1}{n}=\frac{1}{n}$

对于$k>1$，要分两种情况，

对于$i>k$，其最终被选中不仅要求扫描到他时被选中，还要求后来者不被选中或即使选中也不会将它替换。因此对于$j>i$，其不会对$i$构成威胁的概率是：
$F(j)=\frac{k}{j}\times\frac{k-1}{k}+\frac{j-k}{j}=\frac{j-1}{j}$
则$i$被选中的概率：
$\begin{align} P(i) &= \frac{k}{i}\times F(i+1) \times F(i+2) \times\cdots\times F(n) \\\ &= \frac{k}{i} \times \frac{i}{i+1}\times\frac{i+1}{i+2}\times\cdots\times\frac{n-1}{n}\\\ &=\frac{k}{n} \end{align}$
对于$i<=k$，其最终被选中要求第$k$个元素后面的元素不被选中，或者选中但不替换到它。对与$j>k$，其不会对$i$构成威胁的概率等同于$F(j)$。

则$i$被选中的概率为：
$\begin{align} P(i) &=F(k+1)\times F(k+2) \times F(k+3) \times \cdots\times F(n) \\\ &=\frac{k}{k+1}\times\frac{k+1}{k+2}\times\cdots\times\frac{n-1}{n}\\\ &=\frac{k}{n} \end{align}$

证毕

代码如下：

vector<int> reservoirSampling(vector<int> &nums, int k) {
    for (int i = k; i < nums.size(); i++) {
        int r = rand() % (i + 1);
      	// k/i 的概率
        if (r < k) swap(nums[r], nums[i]);
    }
    return vector<int>(nums.begin(), nums.begin() + k);
}

分布式抽样

现在数据量太大，不可能所有数据都放在一台服务器上，如果输入放在$m$台机器上该如何抽样？

不管环境怎么变，问题的本质没有变：产生一个随机的序列，取序列的前$k$个数。最大堆方法可以适用与这种情况。

算法描述：我们在$m$台机器上分别用最大堆方法抽样出最大的$k$个元素，然后将这$mk$个元素发送到一台机器上再用最大堆方法抽样。

通过这种方法，我们得到的还是随机排列中最小的$k$个元素，这是一个随机抽样。

更进一步

如果$k$太大，内存放不下怎么办？

算法描述：可以设置$L$个桶（桶可以放到外存），每次随机的时候不仅产生一个key，同时产生一个桶号$l$，将数据分配到$l$号桶中。默认小号桶比大号桶数据小。统计前$j$个桶的所有元素的个数$S(j)$，如果$j$是第一个$S(j)>k$的桶，在第$j$个桶中抽样$k-S(j-1)$个元素（排个序就行）。因为桶足够多，每个桶中的元素足够少，可以加载到内存。

因为通过排序的关键了$(l,r)$是随机产生的，该排列也是一个随机排列。