核心要点

  • 场景：数据流总量 n 未知或无法一次载入内存，需等概率抽取 k 个样本。

  • 算法：前 k 个直接入池；第 i 个元素（i>k）以 k/i 的概率替换池中随机一个位置。

  • 正确性：任一元素最终留在池中的概率恒为 k/n，与位置无关。

  • 复杂度：单遍扫描 O(n) 时间、O(k) 空间。

标准回答

蓄水池抽样（Reservoir Sampling）解决「流式数据中等概率抽样」问题。先取 k=1 的情形：维护一个变量 res，遍历到第 i 个元素时，以 1/i 的概率把 res 更新为该元素。第 i 个元素最终被保留的概率为：选中它（1/i）× 之后第 i+1..n 个都没替换它，即 (1/i)·(i/(i+1))·((i+1)/(i+2))···((n-1)/n) = 1/n，对所有 i 都成立，故等概率。

一般 k 的情形：前 k 个元素直接放入池中；从第 i=k+1 个元素起，以 k/i 的概率决定是否「入池」，若入池则随机选池中一个位置（0..k-1）替换掉。归纳证明：第 i 个元素被选概率为 k/i（入池），此后每一步存活概率为「不发生替换」+「发生替换但没选到它」= (1 - k/(j) · 1/k) 对 j=i+1..n，乘积化简后整体概率为 k/n。每个元素被选概率相等且为 k/n。

python
import random

def reservoir_sample_one(stream):
    """k=1: 从数据流中等概率抽取 1 个元素"""
    res = None
    for i, x in enumerate(stream, start=1):  # i 从 1 计数
        if random.randint(1, i) == 1:        # 以 1/i 概率替换
            res = x
    return res

def reservoir_sample_k(stream, k):
    """一般 k: 等概率抽取 k 个元素，每个被选概率 k/n"""
    pool = []
    for i, x in enumerate(stream, start=1):
        if i <= k:
            pool.append(x)                   # 前 k 个直接入池
        else:
            j = random.randint(1, i)         # j 在 [1, i]
            if j <= k:                        # 以 k/i 概率入池
                pool[j - 1] = x               # 替换池中随机位置
    return pool

if __name__ == '__main__':
    # 频率验证：每个元素被选概率应约为 k/n
    from collections import Counter
    n, k, trials = 10, 3, 200000
    cnt = Counter()
    for _ in range(trials):
        for v in reservoir_sample_k(range(n), k):
            cnt[v] += 1
    for v in range(n):
        print(v, round(cnt[v] / trials, 3))  # 应都 ≈ k/n = 0.3

常见误区

⚠️ 常见踩坑

不要对第 i 个元素用固定概率（如 1/2）替换；概率必须随 i 递减为 k/i，否则后来的元素会被过度采样，破坏等概率性。

追问

追问 1：为什么 k=1 时第 1 个元素和第 n 个元素被选概率相同？

第 1 个元素以 1/1 入池，但要存活需后续 n-1 步都不替换它，概率为 (1/2)(2/3)···((n-1)/n)=1/n；第 n 个元素直接以 1/n 概率替换入池且无后续，也是 1/n。两者相等。

追问 2：如果要按权重加权采样怎么办？

用 A-Res 算法：为每个元素生成 key = u^(1/w)（u 为 (0,1) 均匀随机，w 为权重），维护 key 最大的 k 个元素（小顶堆）。等价于按权重不放回抽样。

延伸学习

与本题相关的知识库文章、术语、工具与行业资讯。

📖 术语表

🎯 相关面试题