[数据结构与算法]算法学习笔记(1) 蓄水池抽样（在线随机抽样算法）

本文属于「算法学习」系列文章之一。之前的【数据结构和算法设计】系列着重于基础的数据结构和算法设计课程的学习，与之不同的是，这一系列主要用来记录对大学课程范围之外的高级算法学习、优化与使用的过程，同时也将归纳总结出简洁明了的算法模板，以便记忆和运用。在本系列学习文章中，为了透彻讲解算法和代码，本人参考了诸多博客、教程、文档、书籍等资料，由于精力有限，恕不能一一列出。

为了方便在PC上运行调试、分享代码，我还建立了相关的仓库：https://github.com/memcpy0/Algorithm-Templates。在这一仓库中，你可以看到算法文章、模板代码、应用题目等等。由于本系列文章的内容随时可能发生更新变动，欢迎关注和收藏算法学习系列文章目录一文以作备忘。

1. 问题介绍

一个很简单的问题是，从 $n$ 个元素中随机抽取 $k(k\in [1,n])$ 个元素。只是通过放松限制，这一问题也能变得十分棘手。

• 在能将数据全部读入内存的情况下，无论 $n$ 是否事先确定，无论是要随机抽取一个或多个元素，我们都能轻松解决；
• 在不能将数据全部读入内存的情况下，$n$ 是否事先确定、是要随机抽取一个或多个元素，会稍微困难一点；
• 在数据流情况下，由于数据只能被读取一次，且数据量很大、无法全部读入内存，此时完全无法在抽样前确定数据量 $n$ ，但又要随时保证抽取的在线随机 Online Random 性质，就有了这个问题。

怎样才能算是“随机”呢？很简单，就是要使每个数据被抽样的概率都相等。什么是“在线随机”呢？就是在当前来了 $n$ 个数据时选取每个数据的概率均为 $1/n$ ，又来了 $m$ 个数据时选取每个数据的概率均为 $1/(n+m)$ 。

2. 具体解答

让问题简化的一个方法是具体化，比如说在 $100$ 个数据中可以等概率抽取 $10$ 个数据，又来了 $50$ 个数据后，如何等概率抽取 $10$ 个数据？或者看一下《编程珠玑》Column 12中的题目10：在不知道文件总行数的情况下，如何从文件中随机抽取一行？

How could you select one of n objects at random, where you see the objects sequentially but you do not know the value of n beforehand? For concreteness, how would you read a text file, and select and print one random line, when you don’t know the number of lines in advance?

具体来说，在扫描第 $i$ 行时，以 $1/i$ 的概率选择该行，并替换前面已经选到的那行——扫描第一行时，选择第一行；扫描到第二行时，以 $1/2$ 的概率用第二行替换第一行；扫描到第三行时，以 $1/3$ 的概率用第三行替换前面选择的那行……以此类推，直到全部扫描完，此时选到的那一行恰好是等概率的。

用数学来证明，这种策略下选取某行的概率是 $1/n$ 。设行号为 $[1...n]$ ，选到第 $i$ 行（$1\le i\le n$）这一事件的概率等价于，第 $i$ 次选取了第 $i$ 行（并替换了前面选的那行）、且之后每次选择过程中都没能用当前行替换第 $i$ 行。计算结果如下，发现扫描完成后选取某行的概率恰好是 $1/n$ ：
$$\begin{array}{rl}p& =\frac{1}{i}\times (1?\frac{1}{i+1})\times (1?\frac{1}{i+2})\times (1?\frac{1}{i+3})\times ...\times (1?\frac{1}{n})\\ & =\frac{1}{i}\times \frac{i}{i+1}\times \frac{i+1}{i+2}\times ...\times \frac{n?1}{n}=\frac{1}{n}\end{array}$$

在此基础上扩展，现在要随机选择 $k$ 行了。我们的策略是：先选取最前面的 $k$ 行；从第 $k+1$ 行开始，设当前行号为 $i$（$k+1\le i\le n$），以 $k/i$ 的概率选取该行并替换前面选到的 $k$ 行中的一行。

用数学来证明，这种策略下选取某行的概率是 $k/n$ 。设行号为 $[1...n]$ ，选到第 $i$ 行（$k+1\le i\le n$）这一事件的概率等价于，第 $i$ 次选取了第 $i$ 行、并且之后每次选择过程中要么没能用当前行替换第 $i$ 行、要么替换的不是第 $i$ 行。计算结果如下，发现扫描完成后选取某行的概率恰好是 $k/n$ ：
$$\begin{array}{rl}p& =\frac{k}{i}\times (\frac{i+1?k}{i+1}+\frac{k}{i+1}\times \frac{k?1}{k})\times (\frac{i+2?k}{i+2}+\frac{k}{i+2}\times \frac{k?1}{k})\times ...\times (\frac{n?k}{n}+\frac{k}{n}\times \frac{k?1}{k})\\ & =\frac{k}{i}\times \frac{i}{i+1}\times \frac{i+1}{i+2}\times ...\times \frac{n?1}{n}=\frac{k}{n}\end{array}$$

现在可以回答这一问题了——在 $100$ 个数据中可以等概率抽取 $10$ 个数据，又来了 $50$ 个数据后，如何等概率抽取 $10$ 个数据？

• 我们先抽取前 $10$ 个数据，作为备选答案；
• 对后来的第 $i$ 个数据，以 $k/i$ 的概率选择这一数据并替换掉 $10$ 个已选数据中的某个；
• 读取完 $100$ 个数据时，我们手中的 $10$ 个数据是随机等概率（$1/10$）选出的；
• 对后来的 $50$ 个数据采取同样的做法；
• 读取完 $150$ 个数据时，我们手中的 $10$ 个数据依旧是随机等概率（$1/15$）选出的

3. 问题变形和实际应用

结合不同的实际背景，这种题目可能有多种不同的形式，但都可以抽象为蓄水池抽样问题：

• 在一个长度未知的链表中随机选取 $k$ 个元素，要求仅扫描链表一次、且选取 $k$ 个元素是等概率的；
• 从实时的搜素词数据流中随机抽出 $k$ 个词；
• 398. Random Pick Index 给定一个可能含有重复元素的整数数组，要求随机输出（一定存在的）给定数字的索引。蓄水池抽样模板题。
• 528. Random Pick with Weight 给定一个正整数数组 w ，w[i] 代表下标 i 的权重，随机地获取下标 i 并使得选取下标 i 的概率与 w[i] 成正比。暴力使用蓄水池抽样算法可能TLE，要用前缀和+随机化+二分，或者桶轮询+模拟。
• 497. Random Point in Non-overlapping Rectangles 给定一个非重叠轴对齐矩形的列表，随机均匀地选取矩形覆盖的空间中的整数点。暴力使用蓄水池抽样算法可能TLE，需要用前缀和+随机化+二分选择矩形，再在矩形内部用蓄水池抽样算法随机选取整数点。