現(xiàn)在有一組數(shù)，不知道這組數(shù)的總量有多少，請描述一種算法能夠在這組數(shù)據(jù)中隨機抽取k個數(shù)，使得每個數(shù)被取出來的概率相等。

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

蓄水池算法

游泳池（蓄水池）大家都不陌生，有些游泳池中的水是活的，有入水管也有出水管，那么和泳池體積相當?shù)乃鬟^之后，是不是泳池中所有的水都會被替換呢？當然不是，有的水在泳池中可能會存留很久，有的可能剛進去就流走了。仿照這種現(xiàn)象，蓄水池抽樣算法誕生了，蓄水池算法的關(guān)鍵在于保證流入蓄水池的水和已經(jīng)在池中的水以相同的概率留存在蓄水池中。并且蓄水池算法可以在不預(yù)先知道總量的情況下，在時間復(fù)雜度O(N)的情況下，來解決這類采樣問題。

核心原理

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

Python實現(xiàn)

接下來嘗試用Python實現(xiàn)一下蓄水池算法，由于蓄水池算法是在事先不知道總量的情況下抽樣的，所以定義一個方法來接收單個元素，并且把這個方法放在類中，以持有采樣后的數(shù)據(jù)。

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:             self._sample.append(item)             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 復(fù)制代碼

測試代碼

接下來實現(xiàn)一個測試用例驗證實現(xiàn)的算法是否正確，既然是隨機抽樣，無法通過單詞測試來驗證是否正確，所以通過多次執(zhí)行的方式來驗證，比如從1-10里隨機取樣3個數(shù)，然后執(zhí)行10000次取樣，如果算法正確，最后結(jié)果中1-10被取樣的次數(shù)應(yīng)該是相同的，都是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() 復(fù)制代碼

輸出的結(jié)果如下

Counter({7: 3084, 6: 3042, 10: 3033, 3: 3020, 8: 3016, 5: 2997, 4: 2986, 2: 2972, 9: 2932, 1: 2918}) 復(fù)制代碼

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

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

代碼

上述的算法和測試代碼已經(jīng)放在Github，可以直接下載使用。

作者：simpleapples
來源：掘金