[程序员代码面试指南]第9章-蓄水池算法

Posted 2021-12-01 coding-gaga

题目描述

• 从N个元素中随机抽取k个元素，但的k个数无法事先确定。
• 在实际应用中，往往会遇到很大数据流的情况。因此，我们无法先保存整个数据流然后再从中选取，而是期望有一种将数据流遍历一遍就得到所选取的元素，并且保证得到的元素是随机的算法。
• 特别地，此题元素为1-N。

解题思路

• 蓄水池算法。
  1. 先选取个元素中的前k个元素，保存在集合中；
  2. 从第i(i>k)个元素开始，每次先以k/i概率选择是否让第i个元素留下。若第i个元素存活，则从集合中k个元素随机扔掉一个，并将该元素放入集合；否则直接扔掉该元素；
  3. 重复1或2，直到结束。最后集合中剩下的就是保证随机抽取的k个元素。
• 证明：此种方法保证在选第N号球时，从i号球被选中到第N号球的过程中，第i号球最终留在袋子的概率是(k/i)(i/i+1)(i+1/i+2)....(N-1)/N=k/N

代码

public class Main {
    public static void main(String args[]) {
        int k=5,N=1000;
        int[] arr=new int[k];
        arr=getKNumsRandom(k,N);
        for(int i=0;i<k;++i) {
            System.out.println(arr[i]);
        }
    }
    
    public static int[] getKNumsRandom(int k,int N) {
        int arr[]=new int[k];
        for(int i=0;i<k;++i) {
            int num=i+1;
            arr[i]=num;
        }
        for(int i=k;i<N;++i) {
            int num=i+1;//
            if(rand(num)<=k) {//替换
                arr[rand(k)-1]=num;
            }
        }
        return arr;
    }
    
    public static int rand(int max) {
        int test=(int)(Math.random()*max)+1;
        return (int)(Math.random()*max)+1;//随机返回[1,max]的一个值
    }
}