蓄水池算法

 题目：

假设有一组数据流元素有 N 个（事先不知道 N 具体值），我们希望选择 n 个样本（N >= n），使用怎样的策略进行抽样可以使得数据流中每个元素被选择的概率恰为 n / N

结论：

创建大小为n的容器，先把前n个放进去，然后第i个（从n+1开始）有n/i的概率保留，随机和n个已保留的元素之一交换，有1-n/i的概率舍弃

证明：

1.数学归纳法：

        ①当N=n时，每个样本都选择概率都为n/N，显然成立。

        ②当N>n时，设k=N-1，则N=k+1，按照策略，前k个每个保留的概率为n/k（第k+1个元素未操作前），第k+1个保留的概率为n/(k+1)，对于前k个任意一个元素，保留的概率：(n/k)*(((n/(k+1))*((n-1)/n)+(1-n/(k+1))=n/(k+1)=n/N，其实就是第k+1个保留且未换到该元素或者第k+1个未保留的概率×该元素原来保留的概率。

        ③所以当N>=n时，每个样本选择概率都为n/N。

 2.分类推理法：

        按照该策略，对于前n个元素，第i个（i>n）个元素后还保留的概率为(n/i)*((n-1)/n)+(i-n)/i=(i-1)/i

那么到第N个元素还保留的概率：1*(n/(n+1)*((n+1)/(n+2))*...*(N-1)/N=n/N

那么对于第i个元素(i>n)最后保留的概率，(n/i)*(i/(i+1)*...*(N-1)/N=n/N

所以对于所有元素，选择概率都为n/N

 代码实现：

 

import random

def reservoir_sampling(stream, k):
    reservoir = []
    
    # 填充蓄水池，取前k个元素
    for i in range(k):
        reservoir.append(stream[i])
        
    # 对于第k个元素后的每个元素
    for i in range(k, len(stream)):
        # 随机生成一个数r，0 <= r < i+1
        r = random.randint(0, i)
        
        # 如果r小于k，则用当前元素替换蓄水池中的第r个元素
        if r < k:
            reservoir[r] = stream[i]
    
    return reservoir


stream = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
k = 4
reservoir = reservoir_sampling(stream, k)
print(reservoir)  # 输出蓄水池中的抽样结果