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

问题描述

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

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

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

实现方式为如下步骤：

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

代码如下：

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<copy.length; i++) {
            if (copy[i]==target) {
                cnt++;
                int randNum = rand.nextInt(cnt);
                if (randNum<=0) {
                    reservior[0] = i;
                }
            }
        }

        return reservior[0];
    }
}

时间复杂度： $O(N)$ ；空间复杂度： $O(1)$ 。

数学原理

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

对于 $i < k$ 的元素，在 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个元素替代的概率 = $\frac{k}{k+1} \times \frac{1}{k} = \frac{1}{k+1}$
那么反过来第i个元素被保留的概率为 $\frac{k}{k+1}$
那么在 N 步，第 i 个元素被保留的概率应该为：
k+1步被保留的概率 * k+2步被保留的概率 * ... * N步被保留的概率
也就是 $\frac{k}{k+1} \times \frac{k+1}{k+2} \times ... \times \frac{N-1}{N} = \frac{k}{N}$

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

在 k+1步，第k+1个元素被选中的概率为 $\frac{k}{k+1}$ ，由于第 k+1的元素原本不存在，没有先置概率。
在 k+2步，第k+1个元素被保留的概率= 第k+1步被选中概率 * 第k+2步没有选中第k+2个元素的概率
第k+1个元素被保留的概率 = $\frac{k}{k+1} \times \frac{k+1}{k+2} = \frac{k}{k+2}$
在 N 步，第k+1个元素被保留的概率 = $\frac{k}{k+1} \times \frac{k+1}{k+2} \times ... \times \frac{N-1}{N}= \frac{k}{N}$

有几点细节需要留意

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