Introduce#

蓄水池算法 (Reservoir Sampling) 是一种处理流式数据的随机采样问题的高效算法。它可以处理数据量巨大、或者数据动态生成的采样情况，并可以保证每个样本被抽中的概率相等。本章记录了我学习蓄水池算法的一些笔记。

Principle#

蓄水池算法的目标是解决如下的问题：

假设有一个数据集长度为 $N$ ， $N \approx \infty$ ，计算机无法将 $N$ 个元素全部加载都内存中。我们需要从数据集中抽取 $K$ 个元素，要求满足：

每个元素只能访问一次。

每个元素采样概率相等。

蓄水池算法的基本思路是：

维护一个容量为 K 的数组 reservoir[K]，用于存储抽取的元素。
对于第 i 个数据集中的元素，执行如下分类操作：
1. (初始化) 如果 i < K，则将该元素加入采样池 reservoir[i] 中。
2. (替换) 如果 i >= K，则随机生成一个数 r (0 ≤ r ≤ i)，若 r < K，则将 reservoir[r] 替换为当前元素；若 r >= K，则不替换。
抽取结束后，数组 reservoir[K] 为最终的采样结果。

我们通过如下一个例子来解释这个算法的过程：

假设有一个数据集 stream 长度为 $N = 15$ ，数据集为：[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]。我们希望随机采样 $K = 10$ 个元素。

维护一个容量为 K 的数组 reservoir[10]。

对于第 i 个数据集中的元素，执行如下分类操作：
1. (初始化) 如果 i < 10，则将该元素加入采样池 reservoir[i] 中。我们将数据集的前 10 个元素加入 reservoir[0] 到 reservoir[9] 中。
1. (替换) 如果 i >= 10，则随机生成一个数 r (0 ≤ r ≤ i)，若 r < 10，则将 reservoir[r] 替换为当前元素；若 r >= 10，则不替换。
  假设 i = 11, r = 5, r < 10, 则 reservoir[r] = stream[i], 即 reservoir[5] 被替换为 stream[11]。
假设 i = 12, r = 11, r > 10, 则不发生与采样池的替换。
抽取结束后，数组 reservoir 为最终的采样结果。

Code#

实现这样的蓄水池算法很简单，接下来通过一个Python代码来演示。

创建一个函数 reservoir_sampling，接收两个参数 stream 和 K，返回一个长度为 K 的列表 reservoir。
遍历数据集 stream，对每个元素执行分类操作。
1. 如果 i < K，则将该元素加入采样池 reservoir[i] 中。
2. 如果 i >= K，则随机生成一个数 r (0 ≤ r ≤ i)，若 r < K，则将 reservoir[r] 替换为当前元素；若 r >= K，则不替换。
返回最终的采样结果 reservoir。

1
import random
2
def reservoir_sampling(stream, K):
3
    reservoir = []
4

5
    for i, item in enumerate(stream):
6
        if i < K:
7
            reservoir.append(item)
8
            continue
9
        r = random.randint(0, i)
10
        if r < K:
11
            reservoir[r] = item
12

13
    return reservoir

我们可以对当前算法进行一些简单的测试：

1
# 从 100 个元素中随机抽取 10 个
2
stream = [i for i in range(100)]
3
sample_k10 = reservoir_sampling(stream, 10)

Proof#

你可能比较好奇为什么蓄水池算法能保证每个元素被抽中的概率相等。接下来是它的数学证明部分。

我们将其问题转化为数学式：

设数据集的长度为 $N$ ，需要抽取 $K\ (K \le N)$ 个元素，求证：每个元素被抽中的概率为 $P = \frac{K}{N}$ 。

在初始化阶段，第 $i\ (i \le K)$ 个元素一定会被蓄水池采样，此时被抽样的概率 $P_i = 1$ 。在初始化结束后，若 $N = K$ , 则对于第 $i$ 个元素，被抽中的概率为:

\begin{equation} P_{\textbf{i}} = 1 = \frac{K}{K} = \frac{K}{N} \ (1 \le i \le K = N) \end{equation}

根据算法，第 $i \ (i \ge K)$ 个元素会被选入蓄水池的条件是：生成一个 $0$ 到 $i-1$ 的随机数 $r$ ，且 $r < K$ 。

生成随机数的总可能性有 $i$ 种 (从 $0$ 到 $i-1$ ), 使得第 $i$ 个元素被选中的情况有 $K$ 种 ( $0, 1, \dots, K-1$ )。

所以，第 $i$ 个元素被选入蓄水池的概率是：

\begin{equation} P_{\text{select i}} = \frac{K}{i} \end{equation}

选中后，它会随机替换蓄水池中 $K$ 个位置之一。对于任意一个已经在蓄水池中的旧元素 $j$ ，它在处理第 $i-1$ 步时仍在蓄水池中。

那么，元素 $j$ 在第 $i$ 步之后 仍然留在蓄水池 中的概率，可以分为两个事件：

在处理第 $i$ 步之前， $j$ 就在蓄水池中。
在处理第 $i$ 步时， $j$ 没有被替换。

对于事件 1， $j$ 要继续留在蓄水池中，必须满足: 在处理第 $i$ 步时，它没有被第 $i$ 个元素替换。

我们可以使用数学归纳法来继续证明。根据我们的归纳假设，在处理第 $i - 1$ 步时，满足公式 (1)，元素 $j$ 在蓄水池中的概率是：

\begin{equation} P_{(\text{j},\text{i-1})} = \frac{K}{i - 1} \end{equation}

对于事件 2， $j$ 在处理第 $i$ 步时没有被替换，有两种可能，它们之间是或的关系：

第 $i$ 个元素根本没被选中进入蓄水池。这个概率是：1 - P(第 i 个元素被选中)，即：
$\begin{equation} P_{\text{not select i}} = 1 - P_{\text{select i}} = 1 - \frac{K}{i} \end{equation}$
第 $i$ 个元素被选中进入蓄水池，但被蓄水池中的其他元素替换，即 $r \not = {j}$ ，这种情况的概率也很明显是第 $i$ 个元素被选中的概率 $\times$ 没选中 $j$ 的概率。
$\begin{equation} \begin{aligned} P_{\text{not replace j}} &= P_{\text{select i}} \times \frac{K - 1}{K} \\ &=\frac{K}{i} \times \frac{K - 1}{K} \\ &= \frac{K - 1}{i} \end{aligned} \end{equation}$