蓄水池抽樣算法簡介
蓄水池抽樣算法是隨機算法的一種,用來從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個元素開始迭代直到數組結束,算法生成一個隨機數j∈[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. 命題得證。
算法實現
#include<iostream>
#include<vector>
#include<ctime>
using namespace std;
vector<int> ReservoirSampling(vector<int>& v, int n, int k){
if (k > n || v.size() != n)return vector<int>{};
vector<int> res(v.begin(), v.begin() + k);
int i = 0, j = 0;
for (int i = k; i < n; i++){
j = rand() % (i + 1); // inclusive range [0, i]
if (j < k){
res[j] = v[i];
}
}
return res;
}
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;
system("pause");
return 0;
}
結果如圖所示:
算法侷限性
蓄水池算法的基本假設是總的樣本數很多,不能放入內存,暗示了選擇的樣本數k是一個與n無關的常數,然而在實際的應用中,k常常與n相關,比如我們想要隨機選擇1/3的樣本(k = n / 3),這時候需要別的算法或者分佈式算法。