蓄水池采样算法

蓄水池采样算法是一种随机抽样算法，它能够在一个很大的集合中，抽取一部分样本，并保证每个样本的选取概率都是相等并随机的。

• 特点：
  • 选取集合可以非常大，甚至不知道边界。
  • 每个样本的选取随机且概率相等。
  • 时间复杂度较低，O(n)，节省内存。
• 适用场景：
  • 在一些非常大的集合，或者未知大小的集合，不知道边界的集合，不知道文件总行数的情况下，随机抽取k个元素。
  • 保证每个元素抽取都是均匀随机的并且概率相等。
  • 尽量高效，节省内存地抽取

算法原理

蓄水池算法可以抽象为下面的问题:

给定一个数据流，数据流长度N很大，如何从中随机选取k个数据，并且要保证每个数据被抽取到的概率相等。注意：数据流长度N很大，内存无法加载全部数据。

对于该问题, 假设数据序列的规模为 N，需要采样的数量的为 k。

• 首先构建一个可容纳k个元素的数组，然后遍历序列N，对于元素 $N_i$
  • 如果 $i \le k$ ,元素直接放入数组中。
  • 如果 $i \gt k$ ,以 $\frac{k}{i}$ 的概率来决定该元素最后是否被留在数组中(每进来来一个新的元素，数组中的每个旧元素被替换的概率是相同的)。
• 当遍历完所有元素之后，数组中剩下的元素即为所需采取的样本。

证明

对于数组中第i个数据（i≤k）。

• 在 k 步之前，被选中的概率为 1。
• 当走到第 k+1 步时，被第 k+1 个数据替换的概率 = 第k+1个元素被选中的概率 * 第i个数被选中替换的概率，即为:

$$ \frac{k}{k+1} * \frac{1}{k}= \frac{1}{k+1} $$

• 当走到第 k+1 步时，被保留的概率为:

$$ 1 - \frac{k}{k+1} * \frac{1}{k}= \frac{k}{k+1} $$

• 依次类推，在不被第k+1个元素替换的前提下，不被第k+2个数据替换的条件概率为(第k+2个元素被选中的概率*第i个数被选中替换的概率):

$$ 1- \frac{k}{k+2} * \frac{1}{k} = \frac{k+1}{k+2} $$

• 运行到第 n 步时，被保留的概率 = 被选中的概率 * 不被替换的概率，即（条件概率的连乘）：

$$ 1* \frac{k}{k+1} * \frac{k+1}{k+2} * … * \frac{n-1}{n} = \frac{k}{n} $$

对于第j个数据（j>k）。

• 第 j个数据被选中的概率为k/j。
• 被选中的条件下，不被 第j+1 个元素替换的概率为:

$$ 1 - \frac{k}{j+1} * \frac{1}{k}= \frac{j}{j+1} $$

• 运行到第 n步时，被保留的概率 = 被选中的概率 * 不被替换的概率，即条件概率的连乘）：

$$ \frac{k}{j} * \frac{j}{j+1} * \frac{j+1}{j+2} * … * \frac{n-1}{n} = \frac{k}{n}
$$

所以对于其中每个元素，被保留的概率都为 $\frac{k}{n}$.

代码

 public static int[] reservoirSamping(int[] sample, int[] data) {
    // k为样本集合的大小
    int k = sample.length;
    // 从第K+1个元素开始，根据算法随机置换sample中的元素
    for (int i=0; i<data.length; i++) {
        if(i<k){
           // 前k个元素，直接取出，放到sample中
           sample[i] = data[i]; 
        }else{
            // nextInt(i) 返回 [0,i)之间的整数
            // 取得[1, i]之间的随机数，注意Random是左闭右开，所以要加1才能把i包括进去
            int rand = new Random().nextInt(i) + 1;
            // rand小于等于k，则开始置换，将sample中第rand个元素，替换为data中第i个元素
            if (rand <= k) {
                sample[rand-1] = data[i];
            }
        }
    }
    // 遍历一遍，采样完成，返回样本集合sample
    return sample;
}