蓄水池采样

【转】https://zhuanlan.zhihu.com/p/41348264

现在有一组数，不知道这组数的总量有多少，请描述一种算法能够在这组数据中随机抽取k个数，使得每个数被取出来的概率相等。

如果这组数有n个，那么每个数字取到的概率就是k/n，但是这个问题的难点在于不知道这组数的总数，也就是不知道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:
            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()