水塘抽样（应用场景+算法步骤+算法证明+Python实现）

应用场景

在以流形式获取 $n$ 个元素中，随机抽取 $k$ 个样本，满足 $k\le n$（通常有 $k\ll n$）要求：

• 每个元素有相同的概率被选中；
• 时间复杂度不超过 $O(n)$，即遍历一次；空间复杂度不超过 $O(k)$，即要抽取的样本数。

算法步骤

• 初始化长度为 $k$ 的缓存列表 $L$；
• 遍历第 $i$ 个元素：
  • 如果 $i\le k$，则将该元素存入缓存列表；
  • 如果 $i>k$，则在 $[0,1]$ 中取随机数 $t$，若 $t<\frac{k}{i}$，则保留第 $i$ 个元素（保留概率为 $\frac{k}{i}$），并在缓存列表 $L$ 中随机丢掉一个元素（每个元素被丢掉的概率为 $\frac{1}{i}$）；
• 遍历完所有元素后，缓存列表 $L$ 即为抽取的样本。

算法证明

1. 每个元素有相同的概率被选中

证明　使用数学归纳法。当 $n=k$ 时，每个元素被选中的概率均为 $1$，结论显然成立。假设当 $n=m$（$m>=k$）时结论成立，只需证明当 $n=m+1$ 时结论仍成立。

因为当 $n=m$ 时结论成立，所以在遍历第 $m+1$ 个元素前，前 $m$ 个元素被选中的概率相同。因为前 $m$ 个元素被选中的概率相同，所以每个元素被选中的概率为
$$P_m(i)=k÷m=\frac{k}{m}(i=1,2,\cdots ,m)\text{\hspace{1em}(1)}$$
在遍历第 $m+1$ 个元素时，有 $\frac{k}{m+1}$ 的概率会保留第 $m+1$ 个元素，所以第 $m+1$ 个元素被选中的概率为
$$P_{m+1}(m+1)=\frac{k}{m+1}\text{\hspace{1em}(2)}$$
在遍历第 $m+1$ 个元素时，有 $\frac{k}{m+1}$ 的概率会在前 $m$ 个元素选出的 $k$ 个样本中随机丢掉一个，其中每个样本被丢掉的概率为
$$\frac{k}{m+1}\times \frac{1}{k}=\frac{1}{m+1}$$
所以，将式 $(1)$ 代入可得，在遍历第 $m+1$ 个元素后，前 $m$ 个元素被选中的概率为
$$P_{m+1}(i)=P_m(i)\times (1-\frac{1}{m+1})=\frac{k}{m}\times \frac{m}{m+1}=\frac{k}{m+1}(i=1,2,\cdots ,m+1)\text{\hspace{1em}(3)}$$
综上所述，根据式 $(2)$ 和式 $(3)$ 可知，若当 $n=m$（$m>=k$）时结论成立，则当 $n=m+1$ 时，前 $m$ 个元素被选中的概率与第 $m+1$ 个元素被选中的概率相同，均为 $\frac{k}{m+1}$。得证。

2. 时间复杂度不超过 $O(n)$，空间复杂度不超过 $O(k)$。

根据算法步骤显然可知：

• 时间复杂度为 $O(n)$；遍历了所有元素；
• 空间复杂度为 $O(k)$，为缓存列表 $L$ 的长度为 $k$；

满足要求。

代码实现

代码满足 PEP-0008、PEP-0484、pylint 规范；使用 numpy 注释文档规范。

import random
from typing import Any, Iterator, List


def reservoir_sampling(population: Iterator[Any], n_sample: int):
    """水塘抽样

    Parameters
    ----------
    population : Iterator[Any]
        所有元素的可迭代对象（流状态）
    n_sample : int
        需要选取的样本数

    Returns
    -------
    sample : List[Any]
        选中的样本列表
    """
    cache: List[Any] = []  # 缓存列表
    i = 0
    for elem in population:
        if i < n_sample:
            cache.append(elem)
        elif random.random() < n_sample / i:
            cache[random.randint(0, n_sample - 1)] = elem
        i += 1
    return cache

测试用例：

reservoir_sampling(range(5), 5)  # [0, 1, 2, 3, 4]
reservoir_sampling(range(1000), 5)  # [385, 921, 590, 492, 243]