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

如今有一組數，不知道這組數的總量有多少，請描述一種算法可以在這組數據中隨機抽取k個數，使得每一個數被取出來的機率相等。

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

蓄水池算法

游泳池（蓄水池）你們都不陌生，有些游泳池中的水是活的，有入水管也有出水管，那麼和泳池體積至關的水流過以後，是否是泳池中全部的水都會被替換呢？固然不是，有的水在泳池中可能會存留好久，有的可能剛進去就流走了。仿照這種現象，蓄水池抽樣算法誕生了，蓄水池算法的關鍵在於保證流入蓄水池的水和已經在池中的水以相同的機率留存在蓄水池中。而且蓄水池算法能夠在不預先知道總量的狀況下，在時間複雜度O(N)的狀況下，來解決這類採樣問題。

核心原理

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

Python實現

接下來嘗試用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:
            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

測試代碼

接下來實現一個測試用例驗證明現的算法是否正確，既然是隨機抽樣，沒法經過單詞測試來驗證是否正確，因此經過屢次執行的方式來驗證，好比從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({7: 3084, 6: 3042, 10: 3033, 3: 3020, 8: 3016, 5: 2997, 4: 2986, 2: 2972, 9: 2932, 1: 2918})

上面輸出了每一個數字被取樣到的次數，經過圖表能夠清晰的看到分佈狀況

能夠看出蓄水池算法對於隨機抽樣仍是很是適合的，每一個元素的抽樣機率都相同。

代碼

上述的算法和測試代碼已經放在Github，能夠直接下載使用。

關注公衆號【Python私房菜】