蓄水池抽样如何从未知长度数据流中等概率抽取元素？

简化版

蓄水池抽样用于从未知长度的数据流中等概率抽取 k 个元素。抽 1 个时，第 i 个元素以 1/i 的概率替换当前答案。这样处理完全部元素后，每个元素被保留下来的概率都是 1/n。

详细版

抽 1 个元素：

int pick(Iterator<Integer> stream) {
    Random random = new Random();
    int ans = 0;
    int i = 0;
    while (stream.hasNext()) {
        int x = stream.next();
        i++;
        if (random.nextInt(i) == 0) {
            ans = x;
        }
    }
    return ans;
}

抽 k 个时，先放入前 k 个元素；第 i 个元素到来时，在 [0, i-1] 随机选一个位置，如果位置小于 k，就替换蓄水池中的该位置。

完整版教学

一、为什么普通随机下标不适用

如果数组长度已知，可以随机一个下标。但数据流可能很长，甚至不知道什么时候结束，无法提前知道 n。

元素一个个来：x1, x2, x3, ...
不能回头
不能全存
最终还要每个元素等概率

蓄水池抽样就是为这种场景设计的。

二、抽 1 个时的规则

处理第 i 个元素时，以 1/i 的概率选择它替换当前答案。

元素序号	被选为当前答案的概率
第 1 个	`1/1`
第 2 个	`1/2`
第 3 个	`1/3`
第 i 个	`1/i`

看起来后来的元素概率更低，但前面的元素还要经历后续“不被替换”的考验。

关键点：蓄水池抽样不是让每一步概率一样，而是让最终留下来的概率一样。

三、为什么最终概率相等

看第 j 个元素最终留到 n 的概率：

被选中概率 = 1/j
第 j+1 步不被替换 = j/(j+1)
第 j+2 步不被替换 = (j+1)/(j+2)
...
第 n 步不被替换 = (n-1)/n

连乘后中间项抵消：

1/j * j/(j+1) * ... * (n-1)/n = 1/n

所以每个元素最终概率都是 1/n。

四、抽 k 个如何推广

先把前 k 个放入蓄水池。处理第 i 个元素时，随机生成 r，范围 [0, i-1]：

int r = random.nextInt(i);
if (r < k) {
    reservoir[r] = x;
}

第 i 个元素进入蓄水池的概率是 k/i。已有元素被替换出去的概率也会保持平衡。

五、空间优势

无论数据流总长度多大，只保留 k 个元素和计数器。

时间：每个元素处理一次 O(n)
空间：O(k)

这使它适合日志抽样、在线数据分析、大文件随机抽行等场景。

六、随机数边界要写对

Java 的 random.nextInt(i) 生成 0..i-1。抽 1 个时判断 == 0，概率就是 1/i。

if (random.nextInt(i) == 0) ans = x;

如果写成 nextInt(i + 1) 或计数从 0 开始没处理好，概率会偏。

七、常见误区与追问

误区：每个元素都用固定概率替换。 固定概率无法保证最终每个元素等概率。
误区：不知道长度就先全部存下来。 蓄水池抽样的价值就是不需要全存。
误区：随机边界写错。 第 i 个元素应以 1/i 或 k/i 的概率进入。
追问：为什么第一个元素不会概率更大？ 它虽然先被选中，但后续每一步都可能被替换。
追问：抽 k 个的空间是多少？ 只保留蓄水池，空间 O(k)。
追问：适合什么业务场景？ 未知长度流、日志抽样、大文件随机抽取、在线统计。

八、加强记忆

蓄水池抽样的口诀是：第 i 个元素以 1/i 概率替换答案；抽 k 个时以 k/i 概率进入池子。证明靠“先被选中，再一路不被替换”的连乘抵消，最后每个元素都是 1/n。

简化版

详细版