水塘抽样算法的证明

水塘抽样算法主要用来解决这样一类问题：可否在一未知大小的集合中，随机取出一元素？[1] 本文简单对水塘抽样算法进行证明。

水塘抽样算法解决的问题可以形式化地定义为：在大小为 $N$ 的集合 $S$ 中抽取 $k$ 个样本，每个样本被抽中的概率都是 $\frac{k}{N}$，且 $N$ 是未知的。在编程中，这样的问题常见于从一个不知道何时结束的流中随机抽取一些元素，典型的题目如 LeetCode 382. Linked List Random Node 与 398. Random Pick Index。

算法过程

我们定义一个保留 $k$ 个元素的区域叫做保留区。我们以流的形式为例，水塘抽样算法的过程如下：

1. 从流中取得前 $k$ 个元素，填充保留区；
2. 对于第 $i$ 个元素（$k<i\leq N$），我们以 $\frac{k}{i}$ 的概率将其放入保留区，且放入保留区的位置是随机的。

更清楚地讲，放入保留区的位置是随机的即为从保留区中随机确定一个位置替换原有元素，意味着对于当时保留区中的每个元素，被替换的概率是 $\frac{1}{k}$。

证明

按照问题要求，对于流中的每个元素，最终存在于保留区的概率应该为 $\frac{k}{N}$。

设第 $i$ 个元素被放入保留区的概率为 $P_i$，那么其到结束还在保留区中的概率是多少？简单来讲： $$P_i\times(后面的每个元素没被放入保留区的概率+后面每个被放入保留区的元素都没有替换第\ i\ 个元素的概率)$$ 我们一个一个来看。

如果 $i\leq k$

当 $i\leq k$ 时，$P_i=1$。那么对于第 $n$ 个（$i<n\leq N$）元素：

如果 $n\leq k$：

• 没被放入保留区的概率是 $1-0 = 0$，因为一定会被放入；
• 被放入保留区的概率是 $1$，因为一定会被放入；
• 每次有元素被放入保留区，第 $i$ 个元素没有被替换的概率是 $1$，因为在此时不会发生替换。

如果 $n>k$：

• 没被放入保留区的概率是 $1-\frac{k}{n} = \frac{n-k}{k}$；
• 被放入保留区的概率是 $\frac{k}{n}$；
• 每次有元素被放入保留区，第 $i$ 个元素没有被替换的概率是 $1-\frac{1}{k} = \frac{k-1}{k}$。

都写在一起，就得到： $$\begin{align} &P_i \times \left(\prod_{n=i}^{k}(0+1\times 1) + \prod_{n=k+1}^{N}(\frac{n-k}{n}+\frac{k}{n}\times\frac{k-1}{k})\right) \newline =& \prod_{n=k+1}^{N}(\frac{n-k}{n}+\frac{k-1}{n}) \newline =& \prod_{n=k+1}^{N}(\frac{n-1}{n}) \newline =& \frac{k}{k+1}\times\frac{k+1}{k+2}\times…\times\frac{N-2}{N-1}\times\frac{N-1}{N} \newline =& \frac{k}{N} \newline \end{align}$$ 也就是说，第 $i$ 个元素到结束还在保留区中的概率为 $\frac{k}{N}$，符合要求。

如果 $i > k$

当 $i>k$ 时，$P_i=\frac{k}{i}$。对于第 $i$ 个元素后面的第 $n$ 个（$i<n\leq N$）元素：

• 没被放入保留区的概率是 $1-\frac{k}{n} = \frac{n-k}{k}$；
• 被放入保留区的概率是 $\frac{k}{n}$；
• 每次有元素被放入保留区，第 $i$ 个元素没有被替换的概率是 $1-\frac{1}{k} = \frac{k-1}{k}$。

都写在一起，就得到 $$\begin{align} &P_i \times \prod_{n=i+1}^{N}(\frac{n-k}{n}+\frac{k}{n}\times\frac{k-1}{k}) \newline =& \frac{k}{i} \times \prod_{n=i+1}^{N}(\frac{n-k}{n}+\frac{k-1}{n}) \newline =& \frac{k}{i} \times \prod_{n=i+1}^{N}(\frac{n-1}{n}) \newline =& \frac{k}{i} \times (\frac{i}{i+1}\times\frac{i+1}{i+2}\times…\times\frac{N-2}{N-1}\times\frac{N-1}{N}) \newline =& \frac{k}{i} \times \frac{i}{N} \newline =& \frac{k}{N} \newline \end{align}$$ 也就是说，第 $i$ 个元素到结束还在保留区中的概率为 $\frac{k}{N}$，符合要求。

综上，水塘抽样算法得证。

参考文献

1. https://zh.wikipedia.org/wiki/%E6%B0%B4%E5%A1%98%E6%8A%BD%E6%A8%A3