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

题目描述

• 从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]的一个值
	}
}