例题:有一个机器按自然数序列的方式吐出球,1号球,2号球......。现有一个袋子,袋子里最多只能装下k个球,并且除袋子以外没有更多的空间,球扔掉不能放回。设计一种选择方式,使得当机器吐出第N号球时,袋子中的球数是k个,同时可以保证从1号球到N号球中的每一个被选中进袋子的概率都是k/N。
具体例子:
有一个只能装下10个球的袋子:
当吐出100个球时,袋子里有10个球,且1~100号球中每个被选中放入袋子的概率都是10/100;
当吐出1000个球时,袋子里有10个球,且1~1000号球中每个被选中放入袋子的概率都是10/1000;
当吐出N个球时(N>10),袋子里有10个球,且1~N号球中每个被选中放入袋子的概率都是10/N;
蓄水池算法步骤:
1. 处理1~k号球时,直接放入袋子;
2. 处理第i号球(i > k)时,以 k/i 概率决定是否将第i号球放入袋子。若不放入,直接扔掉。若放入,则从袋子里随机选一个球扔掉,然后把第i号球放进袋子。
import random;
class Bag:
selected = []
def __init__(self):
self.selected = []
# 每次拿一个球都会调用这个函数,N表示第i次调用
def carryBalls(self, N, k):
if N <= k:
return self.selected.append(N)
if random.randint(1,N) <= k:
index = random.randint(0,k-1)
self.selected[index] = N
return self.selected