python蓄水池采样算法代码

Python中，蓄水池采样（Reservoir Sampling）算法是一种在线随机抽样方法，用于从大规模数据流中抽取固定大小（n）的样本，而不需要一次性加载所有数据。这个算法非常高效，因为它每次只处理一个元素。

以下是简单的Python代码实现蓄水池采样的步骤：

import random

def reservoir_sampling(stream, n):
    # 初始化一个大小为n的空列表作为蓄水池
    reservoir = [stream[0]] if n > 0 else []
    
    # 遍历数据流中的每个元素，从第i+1个开始
    for i in range(1, len(stream)):
        # 生成一个0到i的随机数
        random_number = random.randint(0, i)
        
        # 如果随机数小于当前蓄水池的大小，将当前元素替换进去
        if random_number < n:
            reservoir[random_number] = stream[i]
            
    return reservoir

# 使用示例
data_stream = list(range(100))  # 假设这是一条很长的数据流
sample_size = 10  # 想要抽取10个样本
random_sample = reservoir_sampling(data_stream, sample_size)
print(random_sample)