问题引入

有一个机器按自然数序列的方式吐出球，1号球，2号球，3号球等等。你没有更多的空间，一个球一旦扔掉，就再也不可拿回。设计一种选择方式，使得当机器吐出第N号球的时候，你袋子中的球数是K个，同时可以保证从1号球到N号球中的每一个，被选进袋子的概率都是K/N。

问题分析

对该问题进行数学建模，即从1-N的自然数序列中等概率抽取K个数，但N是很大或者未知的。对于这种问题最有效的方法就是我们本次要讲的蓄水池采样算法。

蓄水池采样算法

算法描述

1. 使用一个长度为的数组array保存抽取的数字。
2. 对输入序列进行遍历，对于第次遍历，其中，则：
   （1）当 时，全部进行抽取，即。
   （2）当 时，以 概率抽取，并且随机(等概率)替换掉已经抽取了的 个数中其中的一个数。
3. 遍历结束后，数组array中的元素就是最终抽取的数字。

算法证明

• 对于第 个数（）。在步之前，所有数都会被选中，即被选中的概率为 1。当走到第 步时，第个数被第个元素替换的概率为第个元素被选中的概率 * 第被选中替换的概率，即为。则 被保留的概率为 。依此类推，不被第个元素替换的概率为 。则运行到第 步时，第个数被保留的概率 = 被选中的概率 * 不被替换的概率，即：
  

• 对于第个数（）。在第步被选中的概率为。不被个元素替换的概率为 。则运行到第步时，第个数被保留的概率 = 被选中的概率 * 不被替换的概率，即：
  

综上所述，对于每个元素被保留的概率均为 。

算法实现

本文使用java语言实现蓄水池采样算法，如下：

import java.util.*;
public class RadomSelect {
    private int[] selected = null;
    private static Random rand = new Random(12345);

    // 每次拿一个球都会调用这个函数，N表示第i次调用
    public int[] carryBalls(int N, int k) {
        if (selected == null) {
            selected = new int[k];
        }
        if (N <= k) {
            selected[N - 1] = N;
        } else {
            if (rand.nextInt(N) < k) {
                int p = rand.nextInt(k);
                selected[p] = N;
            }
        }
        return selected;
    }

    public static void main(String[] args) {
        Bag bag = new Bag();
        int N = 1000;
        int k = 10;
        for (int i = 1; i <= N; i++) {
            bag.carryBalls(i, k);
        }
        System.out.println(Arrays.toString(bag.selected));
    }
}

总结

蓄水池采样算法适用于在总的样本数未知或者很大的情况下，对样本进行采样的问题。因此，蓄水池采样算法常用于大数据领域。

写在最后

欢迎大家关注我的个人博客复旦猿。