蓄水池抽样算法(Java)

蓄水池抽样算法

蓄水池抽样算法用于解决：在一个数据流的长度很大（大数据内存无法加载），等概率的取出K个数据。

当 k = 1 时， 以1/n选取第n个数。即以我们总是选择第一个对象，1/2的概率选择第二个，以1/3的概率选择第三个…

JAVA代码：

public int[] reservoirSampling(int[] data){
    Random random = new Random();
    int res; 
    int index = 0;
    for (int i : data) {
        if(random.nextInt(++index) == 0){
            res = i;
        }
    }
    return res;
}

对于选取k个对象时，先把读到的前K个对象放入“水库”，从第k+1个对象开始对于第m个对象，以k/m的概率选取，并以1/k的概率替换水库中存在的一个对象，当这个过程结束，每个被选中的概率均为k/n。

JAVA代码：

public int[] reservoirSampling(int[] source, int k){
    Random random = new Random();
    int[] res = new int[k];
    int index = -1;      
    for (int i : source) {
        index++;
        //将前k项先保存进结果数组中
        if (index < k) { 
            res[index] = source[index];
            continue;
        }
        // 从第k+1个对象开始对于第m个对象，以k/m的概率选取，并以1/k的概率替换水库中存在的一个对象
        int r = random.nextInt(index + 1);
        if (r < k) {  
            res[r] = source[index];
        }
    }
    return res;
}