蓄水池抽样算法（Reservoir Sampling）

发表于 | 更新于 | 评论数： | 阅读次数：
本文字数： 2.8k | 阅读时长 ≈ 3 分钟

蓄水池抽样算法（Reservoir Sampling）

许多年以后，当听说蓄水池抽样算法时,将会想起，那个小学数学老师带他做 “小明对水池边加水边放水，求何时能加满水” 应用题的下午。

问题

我是在一次失败的面试经历中听说蓄水池算法的。之后上网搜了搜，知道是一个数据抽样算法，寥寥几行，却暗藏玄机。主要用来解决如下问题。

给定一个数据流，数据流长度 N 很大，且 N 直到处理完所有数据之前都不可知，请问如何在只遍历一遍数据（O(N)）的情况下，能够随机选取出 m 个不重复的数据。

这个场景强调了 3 件事：

  1. 数据流长度 N 很大且不可知，所以不能一次性存入内存。
  2. 时间复杂度为 O(N)。
  3. 随机选取 m 个数，每个数被选中的概率为 m/N。

第 1 点限制了不能直接取 N 内的 m 个随机数，然后按索引取出数据。第 2 点限制了不能先遍历一遍，然后分块存储数据，再随机选取。第 3 点是数据选取绝对随机的保证。讲真，在不知道蓄水池算法前，我想破脑袋也不知道该题做何解。

核心代码及原理

蓄水池抽样算法的核心如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
int[] reservoir = new int[m];

// init
for (int i = 0; i < reservoir.length; i++)
{
    reservoir[i] = dataStream[i];
}

for (int i = m; i < dataStream.length; i++)
{
    // 随机获得一个[0, i]内的随机整数
    int d = rand.nextInt(i + 1);
    // 如果随机整数落在[0, m-1]范围内，则替换蓄水池中的元素
    if (d < m)
    {
        reservoir[d] = dataStream[i];
    }
}

注：这里使用已知长度的数组 dataStream 来表示未知长度的数据流，并假设数据流长度大于蓄水池容量 m。

算法思路大致如下：

  1. 如果接收的数据量小于 m，则依次放入蓄水池。
  2. 当接收到第 i 个数据时，i >= m，在 [0, i] 范围内取以随机数 d，若 d 的落在 [0, m-1] 范围内，则用接收到的第 i 个数据替换蓄水池中的第 d 个数据。
  3. 重复步骤 2。

算法的精妙之处在于：当处理完所有的数据时，蓄水池中的每个数据都是以 m/N 的概率获得的。

下面用白话文推导验证该算法。假设数据开始编号为 1.

第 i 个接收到的数据最后能够留在蓄水池中的概率 = 第 i 个数据进入过蓄水池的概率 * 之后第 i 个数据不被替换的概率（第 i+1 到第 N 次处理数据都不会被替换）。

  1. 当 i<=m 时，数据直接放进蓄水池，所以第 i 个数据进入过蓄水池的概率 = 1。
  2. 当 i>m 时，在 [1,i] 内选取随机数 d，如果 d<=m，则使用第 i 个数据替换蓄水池中第 d 个数据，因此第 i 个数据进入过蓄水池的概率 = m/i。
  3. 当 i<=m 时，程序从接收到第 m+1 个数据时开始执行替换操作，第 m+1 次处理会替换池中数据的为 m/(m+1)，会替换掉第 i 个数据的概率为 1/m，则第 m+1 次处理替换掉第 i 个数据的概率为 (m/(m+1))(1/m)=1/(m+1)，不被替换的概率为 1-1/(m+1)=m/(m+1)。依次，第 m+2 次处理不替换掉第 i 个数据概率为 (m+1)/(m+2)… 第 N 次处理不替换掉第 i 个数据的概率为 (N-1)/N。所以，之后**第 i 个数据不被替换的概率 = m/(m+1)(m+1)/(m+2)…(N-1)/N=m/N**。
  4. 当 i>m 时，程序从接收到第 i+1 个数据时开始有可能替换第 i 个数据。则参考上述第 3 点，之后第 i 个数据不被替换的概率 = i/N。
  5. 结合第 1 点和第 3 点可知，当 i<=m 时，第 i 个接收到的数据最后留在蓄水池中的概率 = 1m/N=m/N。结合第 2 点和第 4 点可知，当 i>m 时，第 i 个接收到的数据留在蓄水池中的概率 = m/ii/N=m/N。综上可知，每个数据最后被选中留在蓄水池中的概率为 m/N。

这个算法建立在统计学基础上，很巧妙地获得了 “m/N” 这个概率。

深入一些——分布式蓄水池抽样（Distributed/Parallel Reservoir Sampling）

一块 CPU 的计算能力再强，也总有内存和磁盘 IO 拖他的后腿。因此为提高数据吞吐量，分布式的硬件搭配软件是现在的主流。

如果遇到超大的数据量，即使是 O(N) 的时间复杂度，蓄水池抽样程序完成抽样任务也将耗时很久。因此分布式的蓄水池抽样算法应运而生。运作原理如下：

  1. 假设有 K 台机器，将大数据集分成 K 个数据流，每台机器使用单机版蓄水池抽样处理一个数据流，抽样 m 个数据，并最后记录处理的数据量为 N1, N2, …, Nk, …, NK(假设 m<Nk)。N1+N2+…+NK=N。
  2. 取 [1, N] 一个随机数 d，若 d<N1，则在第一台机器的蓄水池中等概率不放回地（1/m）选取一个数据；若 N1<=d<(N1+N2)，则在第二台机器的蓄水池中等概率不放回地选取一个数据；一次类推，重复 m 次，则最终从 N 大数据集中选出 m 个数据。

m/N 的概率验证如下：

  1. 第 k 台机器中的蓄水池数据被选取的概率为 m/Nk。
  2. 从第 k 台机器的蓄水池中选取一个数据放进最终蓄水池的概率为 Nk/N。
  3. 第 k 台机器蓄水池的一个数据被选中的概率为 1/m。（不放回选取时等概率的）
  4. 重复 m 次选取，则每个数据被选中的概率为 m(m/NkNk/N*1/m)=m/N。

应用场景

蓄水池抽样的 O(N) 时间复杂度，O(m) 空间复杂度令其适用于对流数据、大数据集的等概率抽样。比如一个大文本数据，随机输出其中的几行。

总结

象征性总结：优雅巧妙的算法——蓄水池抽样。

参考文献

  1. 数据工程师必知算法：蓄水池抽样
  2. 【算法 34】蓄水池抽样算法 (Reservoir Sampling Algorithm)
  3. 分布式 / 并行蓄水池抽样 (Distributed/Parallel Reservoir Sampling)
  4. Distributed/Parallel Reservoir Sampling
多谢支持