力扣中關於蓄水池抽樣問題官方標籤是 2 道，根據我的做題情況來看，可能有三四道。比重算是比較低的，大家可以根據自己的實際情況選擇性掌握。
蓄水池抽樣的算法思維很巧妙，代碼簡單且容易理解，就算不掌握它，作為了解也是很不錯的。
問題描述
給出一個數據流，我們需要在此數據流中隨機選取 k 個數。由於這個數據流的長度很大，因此需要邊遍歷邊處理，而不能將其一次性全部加載到內存。
請寫出一個隨機選擇算法，使得數據流中所有數據被等概率選中。
這種問題的表達形式有很多。比如讓你隨機從一個矩形中抽取 k 個點，隨機從一個單詞列表中抽取 k 個單詞等等，要求你等概率隨機抽取。不管描述怎麼變，其本質上都是一樣的。今天我們就來看看如何做這種題。
算法描述
這個算法叫蓄水池抽樣算法（reservoid sampling）。
其基本思路是：
- 構建一個大小為 k 的數組，將數據流的前 k 個元素放入數組中。
- 對數據流的前 k 個數先不進行任何處理。
- 從數據流的第 k + 1 個數開始，在 [1, i] 之間選一個數 rand，其中 i 表示當前是第幾個數。
- 如果 rand 大於等於 k 什麼都不做
- 如果 rand 小於 k， 將 rand 和 i 交換，也就是説選擇當前的數代替已經被選中的數（備胎）。
- 最終返回倖存的備胎即可
這種算法的核心在於先以某一種概率選取數，並在後續過程以另一種概率換掉之前已經被選中的數。因此實際上每個數被最終選中的概率都是被選中的概率 * 不被替換的概率。
偽代碼：
偽代碼參考的某一本算法書，並略有修改。
Init : a reservoir with the size： k
for i= k+1 to N
if(random(1, i) < k) {
SWAP the Mth value and ith value
}這樣可以保證被選擇的數是等概率的嗎？答案是肯定的。
- 當 i <= k ，i 被選中的概率是 1。
- 到第 k + 1 個數時，第 k + 1 個數被選中的概率（走進上面的 if 分支的概率）是 $\frac{k}{k+1}$，到第 k + 2 個數時，第 k + 2 個數被選中的概率（走進上面的 if 分支的概率）是 $\frac{k}{k+2}$，以此類推。那麼第 n 個數被選中的概率就是 $\frac{k}{n}$
- 上面分析了被選中的概率，接下來分析不被替換的概率。到第 k + 1 個數時，前 k 個數被替換的概率是 $\frac{1}{k}$。到前 k + 2 個數時，第 k + 2 個數被替換的概率是 $\frac{1}{k}$，以此類推。也就是説所有的被替換的概率都是 $\frac{1}{k}$。知道了被替換的概率，那麼不被替換的概率其實就是 1 - 被替換的概率。
因此對於前 k 個數，最終被選擇的概率都是 1 * 不被 k + 1 替換的概率 * 不被 k + 2 替換的概率 * ... 不被 n 替換的概率，即 1 * (1 - 被 k + 1 替換的概率) * (1 - 被 k + 2 替換的概率) * ... (1 - 被 n 替換的概率)，即 $1 \times (1 - \frac{k}{k+1} \times \frac{1}{k}) \times (1 - \frac{k}{k+2} \times \frac{1}{k}) \times ... \times (1 - \frac{k}{n} \times \frac{1}{k}) = \frac{k}{n} $。
對於 第 i (i > k) 個數，最終被選擇的概率是 第 i 步被選中的概率 * 不被第 i + 1 步替換的概率 * ... * 不被第 n 步被替換的概率， 即 $\frac{k}{k+1} \times (1 - \frac{k}{k+2} \times \frac{1}{k}) \times ... \times (1 - \frac{k}{n} \times \frac{1}{k}) = \frac{k}{n} $。
總之，不管是哪個數，被選中的概率都是 $\frac{k}{n}$，滿足概率相等的需求。
相關題目
- 382. 鏈表隨機節點
- 398. 隨機數索引
- 497. 非重疊矩形中的隨機點
總結
蓄水池抽樣算法核心代碼非常簡單。但是卻不容易想到，尤其是之前沒見過的情況下。其核心點在於每個數被最終選中的概率都是被選中的概率 * 不被替換的概率。於是我們可以採取某一種動態手段，使得每一輪都有概率選中和替換一些數字。 上面我們有給出了概率相等的證明過程，大家不妨自己嘗試證明一下。之後結合文末的相關題目練習一下，效果會更好。
