蓄水池抽样算法（Reservoir Sampling）的详细实现与应用

1. 问题描述

蓄水池抽样 是一种在线算法，用于解决一个经典问题：
给定一个数据流（数据长度未知，可能非常大，甚至无法全部放入内存），如何等概率地随机抽取其中 k 个元素？
换句话说，每个元素被选中的概率必须为 \(k / n\)，其中 \(n\) 是数据流最终的总长度（但算法在处理时并不知道 \(n\) 的值）。

2. 核心思想

蓄水池抽样的核心是动态调整概率。
假设我们要维护一个大小为 \(k\) 的“蓄水池”（即结果数组）。算法顺序读取数据流，对于第 \(i\) 个元素（设 \(i\) 从 1 开始计数）：

1. 如果当前已读取的元素数 \(i \leq k\)，则直接放入蓄水池。
2. 如果 \(i > k\)，则以概率 \(k / i\) 决定是否用当前元素替换蓄水池中的某个元素（等概率选择蓄水池中的一个位置进行替换）。
   这样，可以保证在处理完前 \(n\) 个元素后，每个元素留在蓄水池中的概率恰好为 \(k / n\)。

3. 算法步骤（以抽取 k=1 为例，再推广到通用 k）

3.1 最简单情况：k=1（随机选一个元素）

1. 初始化结果变量 reservoir 为 null，计数器 count = 0。
2. 遍历数据流，对于第 i 个元素（i 从 1 开始）：
   • 以概率 \(1 / i\) 将 reservoir 更新为当前元素。
   • 以概率 \(1 - 1/i\) 保持 reservoir 不变。
3. 遍历结束后，reservoir 即为随机选出的元素。

正确性证明（数学归纳法）：

• 当 \(n=1\) 时，概率为 \(1/1 = 1\)，显然成立。
• 假设处理到第 \(n-1\) 个元素时，每个元素被选中的概率为 \(1/(n-1)\)。
• 对于第 \(n\) 个元素：
  • 它被选中的概率是 \(1/n\)。
  • 对于前 \(n-1\) 个元素中的任意一个，它被保留的概率是：

\[ \frac{1}{n-1} \times \left(1 - \frac{1}{n}\right) = \frac{1}{n-1} \times \frac{n-1}{n} = \frac{1}{n} \]

• 因此，每个元素的最终概率都是 \(1/n\)。

3.2 通用情况：任意 k（随机选 k 个元素）

1. 初始化一个大小为 \(k\) 的数组 reservoir。
2. 读取前 \(k\) 个元素，直接放入 reservoir。
3. 从第 \(k+1\) 个元素开始（设当前是第 \(i\) 个元素，\(i > k\)）：
   • 随机生成一个整数 \(r \in [0, i-1]\)（均匀随机）。
   • 如果 \(r < k\)，则用当前元素替换 reservoir[r]。
4. 遍历结束后，reservoir 中即为随机抽取的 \(k\) 个元素。

正确性证明思路：

• 对于前 \(k\) 个元素，它们在最终蓄水池中的概率：
  • 当处理第 \(k+1\) 个元素时，每个前 \(k\) 个元素被替换的概率是：

\[ \frac{k}{k+1} \times \frac{1}{k} = \frac{1}{k+1} \]

因此保留概率为 $ 1 - 1/(k+1) = k/(k+1) $。

• 递推到第 \(n\) 个元素后，每个前 \(k\) 个元素的保留概率为：

\[ \frac{k}{k+1} \times \frac{k+1}{k+2} \times \dots \times \frac{n-1}{n} = \frac{k}{n} \]

• 对于第 \(i\) 个元素（\(i > k\)）：
  • 它被选入蓄水池的概率是 \(k/i\)。
  • 之后不被替换的概率是：

\[ \frac{i}{i+1} \times \frac{i+1}{i+2} \times \dots \times \frac{n-1}{n} = \frac{i}{n} \]

• 因此最终概率为 \((k/i) \times (i/n) = k/n\)。

4. 实现示例（Python）

import random

def reservoir_sampling(stream, k):
    reservoir = []
    for i, item in enumerate(stream):
        if i < k:
            reservoir.append(item)
        else:
            r = random.randint(0, i)  # 生成 [0, i] 之间的随机整数
            if r < k:
                reservoir[r] = item
    return reservoir

注意：这里 random.randint(0, i) 包含两端，所以范围是 \([0, i]\) 共 \(i+1\) 个数，概率 \(k/(i+1)\)。但标准算法中通常用 \(k/i\)，两者在渐进意义上是等价的，只要保证每个位置被选中的概率与 \(i\) 成反比即可。更标准的写法是：

r = random.randrange(0, i+1)  # [0, i]
if r < k:
    reservoir[r] = item

或者用浮点数：

if random.random() < k / (i + 1):
    j = random.randrange(0, k)
    reservoir[j] = item

5. 应用场景

1. 大数据流随机抽样：日志抽样、网络流量分析。
2. 在线广告投放：从实时点击流中抽取样本进行效果分析。
3. 数据库查询优化：从海量数据中快速获取随机样本，用于查询计划评估。
4. 机器学习：在分布式训练中，从数据流中均匀采样。

6. 算法变种与扩展

• 加权蓄水池抽样：每个元素有不同的权重，需要按权重概率抽样。
• 分布式蓄水池抽样：数据分布在多台机器上，每台机器维护一个本地蓄水池，最后合并。
• 滑动窗口蓄水池抽样：只对最近一段时间内的数据流进行等概率抽样。

7. 复杂度分析

• 时间复杂度：\(O(n)\)，只需遍历一次。
• 空间复杂度：\(O(k)\)，只需维护大小为 \(k\) 的蓄水池。
• 随机数生成次数：每处理一个元素（前 \(k\) 个除外）生成一次随机数，共 \(O(n)\) 次。

总结：蓄水池抽样是一种简洁而强大的随机抽样算法，它允许我们在未知总数据量的情况下，仅用一次遍历和少量内存，实现严格的等概率随机抽样。关键在于通过动态调整替换概率，使得每个元素最终被保留的概率均为 \(k/n\)。