【算法34】蓄水池抽樣算法 (Reservoir Sampling Algorithm)

蓄水池抽樣算法簡介

蓄水池抽樣算法隨機算法的一種，用來從 N 個樣本中隨機選擇 K 個樣本，其中 N 非常大（以至於 N 個樣本不能同時放入內存）或者 N 是一個未知數。其時間復雜度為 O(N),包含下列步驟 (假設有一維數組 S, 長度未知，需要從中隨機選擇 k 個元素, 數組下標從 1 開始), 偽代碼如下:

array R[k];    // result
integer i, j;

// fill the reservoir array
for each i in 1 to k do
    R[i] := S[i]
done;

// replace elements with gradually decreasing probability
for each i in k+1 to length(S) do
    j := random(1, i);   // important: inclusive range
    if j <= k then
        R[j] := S[i]
    fi
done

算法首先創建一個長度為 k 的數組（蓄水池）用來存放結果，初始化為 S 的前 k 個元素。然后從 k+1 個元素開始迭代直到數組結束，在 S 的第 i 個元素，算法生成一個隨機數 $j \in [1, i]$， 如果 j <= k， 那么蓄水池的第 j 個元素被替換為 S 的第 i 個元素。

算法的正確性證明

定理：該算法保證每個元素以 k / n 的概率被選入蓄水池數組。

證明：首先，對於任意的 i，第 i 個元素進入蓄水池的概率為 k / i；而在蓄水池內每個元素被替換的概率為 1 / k; 因此在第 i 輪第j個元素被替換的概率為 (k / i ) * (1 / k) = 1 / i。 接下來用數學歸納法來證明，當循環結束時每個元素進入蓄水池的概率為 k / n.

假設在 (i-1) 次迭代后，任意一個元素進入 蓄水池的概率為 k / (i-1)。有上面的結論，在第 i 次迭代時，該元素被替換的概率為 1 / i， 那么其不被替換的概率則為 1 - 1/i = (i-1)/i；在第i 此迭代后，該元素在蓄水池內的概率為 k / (i-1) * (i-1)/i = k / i. 歸納部分結束。

因此當循環結束時，每個元素進入蓄水池的概率為 k / n. 命題得證。

算法的C++實現

實現部分比較簡單，關鍵點也有詳細的注釋，為了驗證算法的正確性，對[1,10]的數組，隨機選擇前五個進行驗證，運行10000次后，每個數字出現的頻率應該是基本相等的，算法的運行結果也證實了這一點，如下圖所示。

#include <iostream>
#include <string>
#include <vector>
#include <cassert>
#include <cstdio>
#include <cstdlib>
#include <ctime>
using namespace std;

/** 
 * Reservoir Sampling Algorithm
 * 
 * Description: Randomly choose k elements from n elements ( n usually is large
 *              enough so that the full n elements may not fill into main memory)
 * Parameters:
 *      v - the original array with n elements
 *      n - the length of v
 *      k - the number of choosen elements
 * 
 * Returns:
 *      An array with k elements choosen from v
 *
 * Assert: 
 *      k <= n
 *
 * Author:  python27
 * Date:    2015/07/14
 */
vector<int> ReservoirSampling(vector<int> v, int n, int k)
{
    assert(v.size() == n && k <= n);

    // init: fill the first k elems into reservoir
    vector<int> reservoirArray(v.begin(), v.begin() + k);

    int i = 0;
    int j = 0;
    // start from the (k+1)th element to replace
    for (i = k; i < n; ++i)
    {
        j = rand() % (i + 1); // inclusive range [0, i]
        if (j < k)
        {
            reservoirArray[j] = v[i];
        }
    }

    return reservoirArray;
}


int main()
{
    vector<int> v(10, 0);
    for (int i = 0; i < 10; ++i) v[i] = i + 1;

    srand((unsigned int)time(NULL));
    // test algorithm RUN_COUNT times
    const int RUN_COUNT = 10000;
    int cnt[11] = {0};
    for (int k = 1; k <= RUN_COUNT; ++k)
    {
        cout << "Running Count " << k << endl;

        vector<int> samples = ReservoirSampling(v, 10, 5);

        for (size_t i = 0; i <samples.size(); ++i)
        {
            cout << samples[i] << " ";
            cnt[samples[i]]++;
        }
        cout << endl;
    }

    // output frequency stats
    cout << "*************************" << endl;
    cout << "Frequency Stats" << endl;
    for (int num = 1; num < 11; ++num)
    {
        cout << num << " : \t" << cnt[num] << endl;
    }
    cout << "*************************" << endl;

    return 0;
}

運行結果如下:

算法的局限性

蓄水池算法的基本假設是總的樣本數很多，不能放入內存，暗示了選擇的樣本數 k 是一個與 n 無關的常數。然而在實際的應用中，k 常常與 n 相關，比如我們想要隨機選擇1/3 的樣本 (k = n / 3)，這時候就需要別的算法或者分布式的算法。

 參考文獻

【1】 Wikipedia：Reservoir Sampling