用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私房菜】