【AI 工业核题 L7】蓄水池抽样（Reservoir Sampling：单遍等概率抽取 k 个样本）（Reservoir Sampling Algorithm）深度实现与原理解析

题目分类：Part L · 经典 ML 与统计模拟 (Part L · Classical ML & Statistical Simulation) | 难度等级：Easy | 工业重要度：工业基石 (核心高频)

一、核心题意与背景

大数据流式处理经典，在总数据量未知甚至无限流中，单遍遍历保证每个样本被选中的概率严格恒为 k/N。

Industrial-grade implementation and mathematical foundations of Reservoir Sampling Algorithm.

二、数学原理与公式推导

数学归纳法严格等概率证明

面对无限网络数据流（总长度 $N$ 无法预先获知且内存装不下）：
1. 将前 $k$ 个元素直接填入容量为 $k$ 的蓄水池；
2. 对于后续到来的第 $i$ 个元素（$i > k$）：
– 以 $frac{k}{i}$ 的概率决定是否收录该元素；
– 若决定收录，在蓄水池已有的 $k$ 个元素中随机挑一个将其替换淘汰。
归纳法证明：
对于任意第 $i$ 个元素，它进入蓄水池的概率为 $frac{k}{i}$。
在后续每个时间步 $j > i$ 中，蓄水池中某个特定位置被新元素替换淘汰的概率为 $frac{k}{j} cdot frac{1}{k} = frac{1}{j}$，因此幸存概率为 $1 – frac{1}{j} = frac{j-1}{j}$。
累乘历史所有幸存概率：
$$P = frac{k}{i} times left(frac{i}{i+1}right) times left(frac{i+1}{i+2}right) dots left(frac{N-1}{N}right) = frac{k}{N}$$
所有中间项分子分母精准对消，证明无论流有多长，每个样本最终留存的概率绝对严格恒等于 $frac{k}{N}$！

三、工业级 Python 核心实现

import numpy as np

def reservoir_sample_stream(stream, k: int) -> list:
    """
    参数:
        stream: 可迭代的无限数据流发生器
        k: 期望等概率抽取的样本数量
    """
    reservoir = []

    for i, item in enumerate(stream):
        if i < k:
            # 前 k 个直接放入蓄水池
            reservoir.append(item)
        else:
            # 依概率 k / (i + 1) 决定是否录取
            j = np.random.randint(0, i + 1)
            if j < k:
                # 替换掉索引为 j 的旧样本
                reservoir[j] = item

    return reservoir

四、自动化单元测试与边界断言

import numpy as np
# 模拟 10000 个流数据抽取 5 个样本
stream = list(range(10000))
sample = reservoir_sample_stream(stream, k=5)
assert len(sample) == 5
assert len(set(sample)) == 5, "样本不应有重复"
print("✓ 蓄水池抽样算法自测通过")

五、张量形状与维度变换流 (Tensor Flow)

• 中文解析：流式输入第 i 个元素 -> 随机整数 j 在 [0, i] -> 若 j < k 则 reservoir[j] = item
• 英文对齐：流式输入第 i 个元素 -> 随机整数 j 在 [0, i] -> 若 j < k 则 reservoir[j] = item

六、工业级数值稳定性避坑清单 (Checklist)

• ⚠️ 生成随机整数的上限必须是 i + 1（即范围为 [0, i]），保证命中概率精确为 k / (i + 1)
• ⚠️ 内存占用极其恒定，永远只驻留 k 个元素

English Checklist:
– Ensure proper multi-dimensional tensor broadcasting and keepdims retention.
– Enforce numerical guards (eps clamping and overflow thresholds) during exponentiation and division.
– Verify train versus eval mode behavioral distinctions (e.g. frozen running statistics and dropout bypass).

七、考场秒记心法口诀

💡 前 k 满仓入，后逐以 k 除以 i 抽签，被抽中者任换池中一员，乘积对消概率齐

Master Reservoir Sampling Algorithm: enforce numerical stability, check tensor shapes, and eliminate redundant memory allocations.

八、高频面试追问与答题策略

Q1：如果流式样本本身带有不同的正权重（Weighted Reservoir Sampling），如何保证每个样本被抽取的概率与其权重正比？
(EN: What are the key trade-offs and memory bottlenecks when deploying Reservoir Sampling Algorithm in high-throughput inference?)

答：使用 Efraimidis-Spirakis 的 A-Res 算法：为每个到来的样本计算一个加权随机键值 $k_i = u_i^{1 / w_i}$（其中 $u_i sim mathrm{Uniform}(0, 1)$），在内存中维护一个容量为 $k$ 的最小堆，保留全局键值最大的前 $k$ 个元素，严格保证采样概率与其权重成正比。

(EN: Memory bandwidth (HBM to SRAM I/O) is the primary latency factor. Fusing element-wise operations and avoiding intermediate tensor materialization significantly outperforms naive implementations.)