【转】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()