不知道蓄水池抽樣算法？那就進來看看吧~

力扣中關於蓄水池抽樣問題官方標籤是 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. 非重疊矩形中的隨機點

總結

蓄水池抽樣算法核心代碼很是簡單。可是卻不容易想到，尤爲是以前沒見過的狀況下。其核心點在於每一個數被最終選中的機率都是被選中的機率 * 不被替換的機率。因而咱們能夠採起某一種動態手段，使得每一輪都有機率選中和替換一些數字。 上面咱們有給出了機率相等的證實過程，你們不妨本身嘗試證實一下。以後結合文末的相關題目練習一下，效果會更好。