蓄水池抽样算法

Posted 2022-07-12 yutianzuijin

最近面试字节，被面试官问到一个这样的问题：给定一个超大的文件，从该文件中随机选取k行文本出来，使得每一行被选择的概率相等。当时一脸懵，后来了解到这个问题需要通过蓄水池抽样算法来解决，经过两天的学习大概搞明白了算法原理，简单记录一下。

1.k=1的情况

k=1时，也就是从文件中选取一行文本，假设总共有n行文本，则必须保证每一行被选择的概率为$1/n$。当我们可以获取所有文本时我们可以很容易利用rand()%n来等概率获取一行文本。但是在很多情况下，我们无法提前预知n的值或者代价较大。比如一个超大文件，无法放入内存，如果想知道行数，我们只能先分块读取整个文件获得行数，然后随机选定一行再重新读一遍文件，这就会导致两次文件读取。又比如数据是流式数据，此时我们根本无法获得n的值。我们希望能有一个算法使得我们在流式读取数据的过程中一遍就可以等概率获得想要的数据，蓄水池算法就是为此而生的。
其基本流程如下：从第一行开始遍历整个文本，当遍历到第i行时，随机选择区间 [0,i)内的一个整数，如果其等于 0，则将答案置为当前行文本，否则答案不变。该算法会保证每一行成为最后被返回的行的概率均为$\frac{1}{n}$，证明如下：
$p(第i行为最后被返回行的概率)$
$=p(第i次随机选择的值为0)\ast p(第i+1次随机选择的值不为0)\ast ...\ast p(第n次随机选择的值不为0)$
$=\frac{1}{i}\ast (1-\frac{1}{i+1})\ast ...\ast (1-\frac{1}{n})$
$=\frac{1}{i}\ast \frac{i}{i+1}\ast ...\ast \frac{n-1}{n}$
$=\frac{1}{n}$

这就证明了按照上述流程获得的文本即为从n行文本中等概率获得的一行文本。同时我们也会发现，按照上述流程我们只需要一次遍历整个文件即可。

2.k>1的情况

当k>1时，也即我们需要从整个文件中等概率选择多行出来。基本思路和一行类似，但是会更为复杂。先描述一下基本操作流程：
对前k行文本，我们需要全部保留；对于第i(i>k)行文本，我们以$\frac{k}{i}$的概率保留，等价描述为随机选择区间$[0,i)$内的一个数，如果小于k则保留；保留下来之后，我们用第i行将之前已保留的k行文本中随机选择的一行进行替换；同样的操作施加到后续的第i+1行直到第n行。按照上述流程选取的k行文本满足每一行被选择的概率为$\frac{k}{n}$，证明如下：

• 对第i行来说，它被选择的概率为$\frac{k}{i}$，无需证明；
• 对第j(0<j<i)行来说，它被选择的概率p为：
  $p(j行保留)=p(j行上一轮被保留)\ast p(i行丢弃)+p(j行上一轮被保留)\ast p(i行保留)\ast p(j行未被替换)$
  $=\frac{k}{i-1}\ast \frac{i-k}{i}+\frac{k}{i-1}\ast \frac{k}{i}\ast \frac{k-1}{k}$
  $=\frac{k}{i}$

从而证明了前i行每一行被选择的概率为$\frac{k}{i}$，从而得证。
此外，当我们证明了k>1的情况再套回k=1的情况时会发现，k=1的情况其实是k>1的特例，只不过在k=1时必定会和前面已选择的一行进行交换而已。