【搞定算法】蓄水池算法

桃扇骨 2021-12-14 13:37 255阅读 0赞

## 1、问题描述分析 ##

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

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

既然说到采样问题，最重要的就是做到公平，也就是保证每个元素被采样到的概率是相同的。所以可以想到要想实现这样的算法，就需要掷骰子，也就是随机数算法。（这里就不具体讨论随机数算法了，假定我们有了一套很成熟的随机数算法了）

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

但是对于第二和第三个问题，就有些不同了，我们不知道数据的整体规模有多大。可能有人会想到，我可以先对数据进行一次遍历，计算出数据的数量 N，然后再按照上述的方法进行采样即可。这当然可以，但是并不好，毕竟这可能需要花上很多时间。也可以尝试估算数据的规模，但是这样得到的采样数据分布可能并不平均。

## 2、蓄水池算法 ##

终于要讲到蓄水池采样算法（Reservoir Sampling）了。先说一下算法的过程：

1、假设数据序列的规模为 n，需要采样的数量的为 k。

2、首先构建一个可容纳 k 个元素的数组，将序列的前 k 个元素放入数组中。

3、然后从第 k+1 个元素开始，以 k/n 的概率来决定该元素最后是否被留在数组中（每进来一个新的元素，数组中的每个旧元素被替换的概率是相同的）。 当遍历完所有元素之后，数组中剩下的元素即为所需采取的样本。

*  ### 证明过程 ###

**第 1 种情况：**对于数组中第 i 个数据（i ≤ k）。在 k 步之前，被选中的概率为 1。当走到第 k+1 步时，被第 k+1 个数据替换的概率 = 第k+1个元素被选中的概率 \* 第i个数 被选中替换的概率，即为![20190704155038700.png][] 。则被保留的概率为![20190704155053676.png][] 。依次类推，在不被第 k + 1 个元素替换的前提下，不被第k+2 个数据替换的条件概率为![20190704155110451.png][]。则运行到第 n 步时，被保留的概率 = 被选中的概率 \* 不被替换的概率，即（条件概率的连乘）：

![20190704155127284.png][]

**第 2 种情况：**对于第 j 个数据（j > k）。第 j个数据被选中的概率为 k / j。不被第 j + 1 个元素替换的概率为![20190704155226221.png][] 。则运行到第 n步时，被保留的概率 = 被选中的概率 \* 不被替换的概率，即条件概率的连乘）：

![20190704155246832.png][]

*  **所以对于其中每个元素，被保留的概率都为 k/n。**

## 3、代码实现 ##

public class ReservoirSampling {
    
        // 从N个元素中等概率的选出K个
        public static int[] sampling(int K, int N){
            if(N < 1 || K < 1 || N < K){
                return null;
            }
            // 初始化所有数据
            int[] arr = new int[N];
            for(int i = 0; i < N; i++){
                arr[i] = i;
            }
    
            int[] pool = new int[K];
            for(int i = 0; i < K; i++){
                // 前K个数据直接放进数组中
                pool[i] = arr[i];
            }
    
            Random random = new Random();
            // K+1个元素开始进行概率抽样
            for(int i = K; i < N; i++){
                // 等概率的返回下标为 0-i中的一个
                int index = random.nextInt(i + 1);
                if(index < K){
                    // 用pool[i]替换掉res[index]，index是随机等概率选中的
                    pool[index] = arr[i];
                }
            }
            return pool;
        }
    
        public static void main(String[] args) {
            System.out.println(Arrays.toString(sampling(20, 10000)));
        }
    }

[20190704155038700.png]: https://img-blog.csdnimg.cn/20190704155038700.png
[20190704155053676.png]: https://img-blog.csdnimg.cn/20190704155053676.png
[20190704155110451.png]: /images/20211214/937644cff5434612a2741f515b09ca89.png
[20190704155127284.png]: /images/20211214/875640ee53f04830bb66b8daf0daf62f.png
[20190704155226221.png]: /images/20211214/1ac4b24c4ef14d11959523b500472db4.png
[20190704155246832.png]: /images/20211214/14a778d2969a4b41ae9a0ec274f237f4.png