9.7蓄水池算法

题目

有一个机器按自然数序列的方式吐出球（1号球，2号球，3号球，······），你有一个袋子，袋子最多只能装下K个球，并且除袋子外，你没有更多的空间。设计一种选择方式，使得当机器吐出第N号球的时候（N>K），你袋子中的球数是K个，同时可以保证从1号球到N号球中的每一个，被选进袋子的概率都是K/N。

核心解法：蓄水池算法

1. 处理1~k号球时，直接放进袋子里。
2. 处理第i号球时（i>k），以k/i的概率决定是否将第i个球放进袋子。如果不决定将第i号球放进袋子，直接扔掉第i号球。如果决定将第i号球放进袋子，那么就从袋子里的k个球中随机扔掉一个，然后把第i号球放入袋子。
3. 处理第i+1号球时重复步骤1或步骤2。

代码实现

public int rand(int max) {
    return ((int) Math.random() * max) + 1;
}

public int[] getKNumsRand(int k, int max) {
    if (max < 1 || k < 1 || k > max) {
        return null;
    }

    int[] res = new int[k];

    for (int i = 0; i != k; i++) {
        res[i] = i + 1; //前k个数直接进袋子
    }

    for (int i = k + 1; i < max + 1; i++) {
        if (rand(i) <= k) { //决定i进不进袋子
            res[rand(k) - 1] = i; //i随机替换掉袋子中的一个
        }
    }

    return res;
}