用Python寫演算法蓄水池演算法實現隨機抽樣

現在有一組數，不知道這組數的總量有多少，請描述一種演算法能夠在這組資料中隨機抽取k個數，使得每個數被取出來的概率相等。

如果這組數有n個，那麼每個數字取到的概率就是k/n，但是這個問題的難點在於不知道這組數的總數，也就是不知道n，那麼該怎麼計算每個數取到的概率呢？

游泳池（蓄水池）大家都不陌生，有些游泳池中的水是活的，有入水管也有出水管，那麼和泳池體積相當的水流過之後，是不是泳池中所有的水都會被替換呢？當然不是，有的水在泳池中可能會存留很久，有的可能剛進去就流走了。仿照這種現象，蓄水池抽樣演算法誕生了，蓄水池演算法的關鍵在於保證流入蓄水池的水和已經在池中的水以相同的概率留存在蓄水池中。並且蓄水池演算法可以在不預先知道總量的情況下，在時間複雜度o(n)的情況下，來解決這類取樣問題。

這一部分涉及公式，為了保證效果直接貼了圖過來。

接下來嘗試用python實現一下蓄水池演算法，由於蓄水池演算法是在事先不知道總量的情況下抽樣的，所以定義乙個方法來接收單個元素，並且把這個方法放在類中，以持有取樣後的資料。

import random
class reservoirsample(object):
def __init__(self, size):
self._size = size
self._counter = 0
self._sample = 
def feed(self, item):
self._counter += 1
# 第i個元素（i <= k），直接進入池中
if len(self._sample) < self._size:
return self._sample
# 第i個元素（i > k），以k / i的概率進入池中
rand_int = random.randint(1, self._counter)
if rand_int <= self._size:
self._sample[rand_int - 1] = item
return self._sample

接下來實現乙個測試用例驗證實現的演算法是否正確，既然是隨機抽樣，無法通過單詞測試來驗證是否正確，所以通過多次執行的方式來驗證，比如從1-10裡隨機取樣3個數，然後執行10000次取樣，如果演算法正確，最後結果中1-10被取樣的次數應該是相同的，都是3000上下。

import unittest
from collections import counter
from reservoir_sample import reservoirsample
class testmain(unittest.testcase):
def test_reservoir_sample(self):
samples = 
for i in range(10000):
sample = 
rs = reservoirsample(3)
for item in range(1, 11):
sample = rs.feed(item)
samples.extend(sample)
r = counter(samples)
print(r)
if __name__ == '__main__':
unittest.main()

輸出的結果如下

counter()

上面輸出了每個數字被取樣到的次數，通過圖表可以清晰的看到分布情況

可以看出蓄水池演算法對於隨機抽樣還是非常適合的，每個元素的抽樣概率都相同。

用Python寫演算法 蓄水池演算法實現隨機抽樣

蓄水池演算法

蓄水池演算法

蓄水池演算法

相關推薦

用Python寫演算法蓄水池演算法實現隨機抽樣