蓄水池抽样算法

问题描述

问题：有一段很长整数流，不知道其长度，内存可能也不够存下所有数据。希望从这个整数流中抽样出 $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)$ 是随机产生的，该排列也是一个随机排列。