蓄水池采样

蓄水池采样（Reservoir Sampling）是一种用于从一个未知大小的数据流中随机选择 k 个样本的算法。

目的：在不需要预先知道数据流总长度 n 的情况下，从数据流中等概率地抽取 m 个元素。
等概率要求：最终结果集（蓄水池）中，任意一个元素 A[k] 被抽中的概率必须是。

算法步骤

初始化一个大小为 k 的数组（蓄水池）来存储样本。
读取数据流的前 k 个元素，直接将它们放入蓄水池中。
对于数据流中的第 i 个元素（i 从 k+1 开始）：
- 以概率 k/i 决定是否将该元素加入蓄水池。
- 如果决定加入，则随机选择蓄水池中的一个元素进行替换。
- 否则，跳过该元素。
重复步骤 3，直到数据流结束。
返回蓄水池中的 k 个样本。

代码实现

#include <iostream>
#include <vector>
#include <cstdlib>
#include <ctime>
using namespace std;

vector<int> reservoirSampling(const vector<int>& nums, int k) {
    vector<int> reservoir;
    int n = 0; // 记录数据流中的元素数量， 即当前处理到第几个元素

    // 初始化随机数生成器
    srand(time(0));

    for (int value : nums) {
        n++;
        if (n <= k) {
            reservoir.push_back(value);  // 直接加入前 k 个
        } else {
            int j = rand() % n;          // 生成 [0, n-1] 的随机数
            if (j < k) {                 // 以 k/n 的概率替换, j 是均匀随机取值于 [0, n-1],  这 n 个数中只有前 k 个满足 j < k, 也就是说取到这 k 个数的概率是 k/n
                reservoir[j] = value;  // 替换蓄水池中的一个元素, 具体替换哪个元素已经由 j 决定
            }
        }
    }

    return reservoir;
}

原理证明

概率证明

假设我们正在处理数据流中第 i 个元素 A[i]（i > m）：

概率决策：以的概率决定是否将 A[i] 放入蓄水池。
替换操作：如果决定放入，则随机选择蓄水池中任意一个位置进行替换。

为什么 A[k] 被选中的概率是？

我们将证明，对于数据流中的任意一个元素 A[k]（1 ≤ k ≤ n），它在整个算法结束后仍然留在蓄水池中的概率是。

1. A[k] 是前 m 个元素（k ≤ m）

A[k] 初始被选中的概率是 1（因为它在前 m 个元素中）。
关注它在后续遍历中不被替换的概率。

对于第 i 个元素 A[i]（i > m）：

A[i] 被选入的概率是。
如果被选入，随机替换蓄水池 m 个位置中的一个，A[k] 被替换的概率是。
所以 A[k] 在第 i 步被替换的概率为： $被替换$
不被替换的概率为： $不被替换$

A[k] 最终保留在蓄水池中的概率是： $最终保留$

这是一个伸缩积（Telescoping Product）： $最终保留$

2. A[k] 是后 n − m 个元素（k > m）

A[k] 在第 k 步被选中的概率是。
在后续 i 步（i > k）不被替换的概率为。

所以 A[k] 最终保留的概率为： $最终保留$

同样是伸缩积，化简后： $最终保留$

总结

无论是前 m 个元素还是后 n − m 个元素，它们最终保留在蓄水池中的概率都是严格相等的。这证明了的替换概率策略能够实现等概率抽样。

应用场景

大数据处理：在处理大规模数据集时，蓄水池采样可以有效地从数据流中抽取样本，而不需要存储整个数据集。
在线算法：适用于需要实时处理数据流并进行随机采样的场景。
机器学习：在训练模型时，可以使用蓄水池采样从大量数据中抽取代表性样本。
统计分析：用于从连续的数据流中获取随机样本以进行统计分析。