蓄水池算法

要解決的問題

假設有一個源源吐出不同球的機器, 只有裝下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. 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;
        }

    }
}

更多

算法和數據結構筆記

參考資料