蓄水池算法

要解决的问题

假设有一个源源吐出不同球的机器， 只有装下10个球的袋子，每一个吐出的球，要么放入袋子，要么永远扔掉,如何做到机器吐出每一个球之后，所有吐出的球都等概率被放进袋子里

规则

吐出1到10号球，完全入袋, 引入随机函数f(i)，提供一个值i，等概率返回1-i的一个数字, 当K号球吐出的时候（K>10) ,我们通过以下决策决定是否要入袋

1. 引入随机函数：f(K) , 如果返回10以内的数，则入袋，如果返回10以外的数，则扔掉, 即：10/K的概率决定球是否入袋。

2. 第一步中如果决定入袋，那么袋子中已经存在的球以等概率丢弃一个。

证明

情况1

当K为1～10号的时候，根据我们的规则，入袋概率100%，每个球等概率

情况2

当K为任意一个大于10的数，我们假设K为927，即：当927这个编号的球吐出来的时候，我们考虑：

A. 1 ～ 10 号球的入袋概率

B. 大于10号小于等于927号球的入袋概率

如果A和B都可以说明等概率

那么可以推广到普遍情况都是等概率的。

A情况，927号球到来的时候，我们可以考虑一下5号球的入袋概率是多少？

5号球需要存活到927号球到来，必须满足：

1. 11号球来的时候，5号球活下来

2. 12号球来的时候，5号球活下来

3. ...

4. 926号球来的时候，5号球活下来

11号球来的时候，5号球怎么活下来呢？先看一下，11号球到来，5号球如果没有活下来的概率q，那么5号球活下来的概率就是 1 - q

首先，根据我们的规则，11号球要以10/11概率被选中入袋，且5号球要以非常倒霉的情况1/10概率被选中要替换，那么，11号球到来，5号球被替换的概率为：

(10/11 * 1/10) = 1/11

那么5号球活下来的概率就是

1 - 1/11 = 10/11

12号球到来的时候，5号球活下来的概率同理，可以计算出为11/12

13号球到来的时候，5号球活下来的概率同理，可以计算出为12/13

....

927号球到来的时候，5号球活来的概率同理：可以计算出为926/927

所以，5号球活下来的概率为：

10/11 * 11/12 * 12/13 ... * 925/926 * 926/927 = 10/927

同理，1 ～ 10 号球任意一号都可以按照5号球的计算方式计算，概率均为：10/927

A情况是等概率的

再看B情况，我们可以假设一个大于10但是小于927的球，比如，15号球，考虑入袋概率

15号球要在927号球到来的时候，还在袋子中，则需要保证：

15号球在当时被吐出的时候，以10/15概率被选中，且在

16号球到来的时候，15号球活下来，概率根据A的计算逻辑，为15/16

17号球到来的时候，15号球活下来, 同理，为16/17

....

926号球到来的时候，15号球活下来，概率为925/926

927号球到来的时候，15号球活下来，概率为926/927

所以，927号球到来的时候，15号球活下来的概率是：

10/15 * 15/16 * 16/17 .... * 925/926 * 926/927 = 10/927

同理，任意大于10小于927号球的概率都可以根据15号球的计算逻辑推算出来，均为10/927

A情况和B情况概率都是10/927

所以规则满足了题目的要求。

代码

public class Code_0058_ReservoirSampling {
    public static class RandomBox {
        private int[] bag;
        // 袋子容量
        private int capacity;
        // 第几号球
        private int count;

        public RandomBox(int capacity) {
            bag = new int[capacity];
            this.capacity = capacity;
            count = 0;
        }

        // 随机函数，等概率生成1-max之间随机的一个数字
        // Math.random() -> 生成[0,1)范围内的数
        // (int)i 是对i进行向下取整
        private int rand(int max) {
            return (int) (Math.random() * max) + 1;
        }

        public void add(int num) {
            // 球个数增加
            count++;
            // 如果球的个数没有超过容量
            if (count <= capacity) {
                // 则入袋
                bag[count - 1] = num;
            } else if (rand(count) <= capacity) {
                // 否则以N/count的概率入袋
                bag[rand(capacity) - 1] = num;
            }
        }

        // 返回袋子中最终选中的球
        public int[] choices() {
            int[] res = new int[capacity];
            System.arraycopy(bag, 0, res, 0, capacity);
            return res;
        }

    }
}

更多

算法和数据结构笔记

参考资料

• 程序员代码面试指南（第2版）
• 算法和数据结构体系班-左程云