蓄水池算法python

蓄水池算法是一种常见的抽样算法，用于从一个包含未知大小的数据流中随机抽取k个样本，保证每个元素被选的概率相等。

在Python中，可以通过以下步骤实现蓄水池算法：

1. 从数据流中读取前k个元素，存储在一个列表中；

2. 遍历流中剩余的元素，对于每个元素，有1/k的概率被选择，即生成一个0到1之间的随机数R，当R小于1/k时，将该元素替换掉所存储的列表中的一个元素；

3. 遍历完成后，存储的列表即为抽取的k个样本。

以下是一个Python示例代码：

import random

def reservoir_sampling(stream, k):
    reservoir = []  # 存储k个样本的列表
    for i, element in enumerate(stream):
        if i < k:
            reservoir.append(element)
        else:
            # 以1/i的概率替换掉列表中的一个元素
            r = random.randint(0, i)
            if r < k:
                reservoir[r] = element
    return reservoir

该示例代码中，stream代表数据流，k代表需要抽取的样本数量。当数据流与k均较大时，蓄水池算法能够在常数时间复杂度下完成抽样，具有高效的特点。