[LeetCode] 算法 - 蓄水池

问题

采样问题经常会被遇到，比如：

• 从 100000 份调查报告中抽取 1000 份进行统计。
• 从一本很厚的电话簿中抽取 1000 人进行姓氏统计。
• 从 Google 搜索 "Ken Thompson"，从中抽取 100 个结果查看哪些是今年的。

这些都是很基本的采用问题。而采样问题，最重要的就是做到公平，也就是保证每个元素被采样到的概率是相同的。

对于第一个问题，还是比较简单，通过算法生成 [0,100000−1) 间的随机数 1000 个，并且保证不重复即可。

但是对于第二和第三个问题，我们不知道数据的整体规模有多大。

算法

蓄水池就解决了这类问题。算法的过程：

假设数据序列的规模为N，需要采样的数量的为K。

• 首先构建一个可容纳K个元素的数组，将序列的前K个元素放入数组中。
• 然后从第K+1个元素开始，以K/i的概率来决定该元素是否被替换到数组中（数组中的元素被替换的概率是相同的）。
• 当遍历完所有元素之后，数组中剩下的元素即为所需采取的样本。

证明

1）当i <= K时，就是前K个元素，最开始都是在池里的。要想最后还能留在池里，就要保证，之后的元素没有替换掉它。这个概率是：P(K+1没替换它) * P(K+2没替换它) *...* P(N没替换它) = (1 - P(K+1替换它)) * (1 - P(K+1替换它)) *...* (1 - P(N替换它))。
P(K+1替换它) = K/(K+1) => P(K+1替换它) = 1 - K/(K+1) = 1/(K+1)。
所以，最后留下来的概率是：K/N。
2）当i > K时，就是K后面的元素，最开始都不在池里。要想最后还能留在池里，就要保证，i被选中加入池里，并且i之后的元素没有替换掉它。这个概率是：P(i被选中加入池中) * P(i+1没替换它) *...* P(N没替换它) = K/i * (1 - P(i+1替换它)) *...* (1 - P(N替换它))。
P(i+1替换它) = K/(i+1) * 1/K => P(i+1替换它) = 1 - 1/(i+1) = i/(i+1)。
所以，最后留下来的概率是：K/N。

代码

public class ReservoirSampling {

    private int[] pool; // 所有数据
    private Random random = new Random();

    public ReservoirSampling(int N) {
        pool = new int[N];
        for (int i = 0; i < N; i++) {
            pool[i] = i;
        }
    }

    private int[] sampling(int K) {
        int[] result = new int[K];
        for (int i = 0; i < K; i++) { // 前 K 个元素直接放入数组中
            result[i] = pool[i];
        }

        for (int i = K; i < N; i++) { // K + 1 个元素开始进行概率采样
            int r = random.nextInt(i + 1);
            if (r < K) {
                result[r] = pool[i];
            }
        }

        return result;
    }
}

相关题目

382 Linked List Random Node
https://leetcode.com/problems/linked-list-random-node/
蓄水池问题，K=1

398 Random Pick Index
https://leetcode.com/problems/random-pick-index/
蓄水池问题，K=1