[算法总结] - 蓄水池采样算法

问题描述

在长度为N的数组中，随机等概率选取K个元素，如何实现这个随机算法。 思路很简单，生成一个[0, N]的随机数index，然后返回index上的数值即可。

但是，如果输入是一个长度未知的数组比如stream，先遍历得到数组大小，在遍历进行K次采样显然不够高效，这就引出了蓄水池算法。

• 蓄水池采样算法可以在一次遍历中得到K次采样结果并且保证等概率
• N个样本 K次采样每一个元素被pick的概率是 k/N

实现方式为如下步骤：

1. 构建一个长度为K的数组(蓄水池)，保存采样结果
2. 将数组[0, k]数值，赋值给蓄水池数组
3. 遍历剩下[k+1, N]，每一次迭代中产生一个 的index, 如果index < K那么将原来处在该index的结果覆盖掉。以此类推
4. 最后返回蓄水池数组结果

代码如下：

Leetcode 398. random pick index

class Solution {

    int[] reservior;
    Random rand = new Random();
    int[] copy;
    public Solution(int[] nums) {
        // 本题目只需要选取一个样本 k = 1
        copy = nums;
        reservior = new int[1];
        reservior[0] = -1;
    }
    
    public int pick(int target) {
        int cnt = 0;
        for (int i=0; i

时间复杂度：；空间复杂度：。

数学原理

上述步骤中最难理解无非就是第三步，为什么这样做就可以实现每一个元素被选的概率是k/N。

对于  的元素， 在 k 步之前，他们被选中是没有随机性的 p = 100%；

• 在 k+1 步时，被第k+1个元素替代的概率 = (k+1)元素被选中的概率 * i 这个index被选中的概率，根据上面实现，第 i 个index被选中概率为 1/k (Java中random.nextInt是左闭右开)，而 k+1个元素被选中的概率为 k/k+1(random生成的随机数小于k都为选中) 
  • 被第k+1个元素替代的概率 = 
  • 那么反过来第i个元素被保留的概率为 
• 那么在 N 步，第 i 个元素被保留的概率应该为：
  • k+1步被保留的概率 * k+2步被保留的概率 * ... * N步被保留的概率
  • 也就是  

对于  的元素，在k步之前，是没有概率的因为不存在

• 在 k+1步，第k+1个元素被选中的概率为  ，由于第 k+1的元素原本不存在，没有先置概率。
• 在 k+2步，第k+1个元素被保留的概率= 第k+1步被选中概率 * 第k+2步没有选中第k+2个元素的概率
  • 第k+1个元素被保留的概率 = 
• 在 N 步，第k+1个元素被保留的概率 = 

有几点细节需要留意

1. 所有的数值，只有一次选中的机会，就是数组遍历到那个index的时候，如果没有被选中，那么以后再也没有机会被重新选中。只有当时被选中才有保留的机会 
   1. [0, k]的元素第一次被选中概率为 100%
   2. [k+1, N]的元素第一次被选中概率为  
2. 不管数组中那个元素只要被选中，保留到最后作为返回值的概率都是