[杂谈] 16. Reservoir Sampling 蓄水池抽样算法

1. 常见例题

采样问题常见面试题：

• 有一堆数据，但是你不知道有多少个，让你随机抽取10个？
• 从 100000 份调查报告中抽取 1000 份进行统计。
• 从一本很厚的电话簿中抽取 1000 人进行姓氏统计。
• 从 Google 搜索 “Ken Thompson”，从中抽取 100 个结果查看哪些是今年的。

《编程珠玑》中的一个习题：

• 如何随机从n个对象中（这n个对象是按序排列的，但是在此之前你是不知道n的值的）随机选择一个对象？

相关OJ:

• LeetCode 382. 链表随机节点
• LeetCode 398. 随机数索引

假设要从一大堆元素中选择k个元素，不知道多少，大数据。蓄水池抽样算法给了我们在不知道总数的情况下，等概率随机抽样的方法。

2. 算法的正确性证明

使用蓄水池抽样算法的具体的思路是：先初始化一个集合，集合中有k个元素，将此集合作为蓄水池。然后从第k+1个元素开始遍历，并且按一定的概率替换掉蓄水池里面的元素。

用途：该算法保证每个元素以 k / n 的概率被选入蓄水池数组。

取前 k个元素放入蓄水池中。从 i = k + 1开始，以 k／i 的概率取第 i 个元素。若第 i 个元素被选中，已均等的概率(即 1 / k )替换蓄水池中的先前被选中的任一元素。

这样，我们就证明了用蓄水池抽样算法抽取每个元素的概率是相等的。

那么如果了解了蓄水池抽样，LeetCode 398. 随机数索引这道题就不算一道难题了。定义两个变量，计数器 cnt 和返回结果 res，遍历整个数组，如果数组的值不等于 target，直接跳过；如果等于 target，计数器加 1，然后在 [0,cnt) 范围内随机生成一个数字，如果这个数字是 0，将 res赋值为i即可。

参见代码如下：

class Solution {
public:
    Solution(vector<int> nums): v(nums) {}
    
    int pick(int target) {
        int cnt = 0, res = -1;
        for (int i = 0; i < v.size(); ++i) {
            if (v[i] != target) continue;
            ++cnt;
            if (rand() % cnt == 0) res = i;
        }
        return res;
    }
private:
    vector<int> v;
};