【算法】蓄水池抽样

版权声明：转载请注明 https://blog.csdn.net/u013166817/article/details/84501618

例题：有一个机器按自然数序列的方式吐出球，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