蓄水池抽样问题

发表于 | 分类于 |

前几天听到同事面试的时候问了蓄水池抽样问题，以前也听过这个问题，但是也没有自习思考过这个问题，一时间也想得不是很清楚这个问题，就花了一点时间看了一些，这里做一个记录，也让自己以后想起的时候有一个回忆的地方。

什么是蓄水池抽样问题呢？

先来看一下维基百科的解释：

Reservoir sampling is a family of randomized algorithms for randomly choosing a sample of k items from a list S containing n items, where n is either a very large or unknown number. Typically n is large enough that the list doesn’t fit into main memory.

也就是从n个元素中随机地选择出k个元素，需要满足每个元素被选择出来的概率是一样的，同时这里的n预先是未知的，每一次接收输入的时候选择的时候都需要使每个数被选择出来的概率相等。

分析

最简单的方法就是当每一次接收到一个数的时候，我们从当前接收到的所有树中重新选择k个数出来，每次选择的时候随机选择也就达到了等概率，但是显而易见的问题就是浪费资源，每次选择的话复杂度有点高，我们需要一个O(n)的算法来完成。下面来解决这个问题，我们考虑两种情况，k等于1和k不等于1.

k=1

k=1是一种特殊情况，在这种情况下，我们只需要每次接收到输入的时候，保证每个数被选中的概率是1/k即可。证明用数学归纳法即可，大概如下如下：

  1. 假设第m个数读入进来之后，第m个数被选中的概率是1/m
  2. 当m=m+1的时候，第m+1个数被选中的概率是1/(m+1)，前m个数被选中的概率是1/m * (m/(m+1))=1/(m+1)

具体实现代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
void reservoir_sampling_single()
{
    int res = -1, tmp = 0;
    int num = 0;
    while (cin >> tmp)
    {
        if (rand() % (num + 1) == num)
        {
            res = tmp;
        }
        num ++;
    }
    cout << "the random number is: " << res << endl;
}

注：这里代码也有不严谨的地方那个，如res=-1这句，为了方便也不再进行处理。

k!=1

这个时候只要保证每个数被选中的概率是k/n即可。这里证明同上，将1换成k即可。

实现代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
void reservoir_sampling_multi(int k)
{
    int *res = new int[k];
    int num = 0;
    int tmp = 0;
    while (cin >> tmp)
    {
        if (num < k)
        {
            res[num ++] = tmp;
            continue;
        }
        if (rand() % (num + 1) < k)
        {
            res[rand() % k] = tmp;
        }
        num ++;
    }
    for (int i = 0; i < k; ++i) {
        cout << "the random nums are: " << res[i] << endl;
    }
}

后记

算法是挺有意思的东西，有时间多思考思考挺好，之前本来想详细谢谢败者树的，最近时间也是不允许，根据工作需要，可能会详细弄弄推荐相关的内容，也期待在这里总结分享。

参考文献

Reservoir sampling