蓄水池抽样算法

问题描述

给出一个数据流，这个数据流的长度很大或者未知（内存无法一次性容纳下），并且对该数据流中数据只能访问一次。
简而言之：要求从N个元素中随机的抽取k个元素，其中N的大小未知，k >= 1。
请写出一个随机选择算法，使得数据流中所有数据被选中的概率相等。

算法思路

先初始化一个集合，集合中有k个元素，将此集合作为蓄水池（reservoir），然后从第k+1个元素开始遍历，并且按一定的概率替换掉蓄水池里面的元素。

《The Art of Computer Programming》中的伪代码

init : a reservoir with the size： k
for i= k+1 to N  
    M = random(1, i);  
    if( M < k)  
    SWAP the Mth value and ith value  
end for   

先将前k个数取出来放入蓄水池中，然后从第k+1个数开始遍历。假设遍历到第i个数，以k/i的概率替换掉蓄水池中的某个元素即可。

数学归纳法证明

假设i=n时，前k个元素都以k/n被选中；
那么当i=n+1是，第n+1个元素被选中的概率为k/n+1；
对于前面的n个元素，每个元素被选中的情况分为两种：
1.前面n次已经被选中，第n+1次时，第n+1个元素没有被选中；
2.前面n次已经被选中，第n+1次时，第n+1个元素被选中但是没有将其替换掉;
此时的概率为： k/n×(1−k/n+1)+k/n×(k/n+1×(1−1/k))=k/n+1
由此可见，第n+1步也满足假设条件，问题得到证明。

Java代码实现

    public static List reservior(int k){
        List raw = getRandomList(1, 100000);
        List res = new ArrayList(k);
        for(int i = 0; i < k; i++){
            res.add(raw.get(i));
        }
        for(int i = k; i < raw.size(); i++){
            int m = rand(0,i);
            if(m < k){
                swap(res.get(m), raw.get(i));
            }
        }
        return res;
    }

    private static List getRandomList(int start, int end) {
        List res = new ArrayList();
        for (int i = start; i <= end; i++) {
            res.add(i);
        }
        for (int i = 0; i < lens; i++) {
            int t = rand(i, lens);
            swap(res.get(i), res.get(t));
        }
        return res;
    }