用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 上下。 ...