Reservoir Sampling 蓄水池采样算法

本文最后更新于：2023年12月5日下午

长数据流的随机采样可以使用蓄水池采样算法，本文记录相关内容。

简介

问题描述：给定一串很长的数据流，对该数据流中数据只能访问一次，使得数据流中所有数据被选中的概率相等。

解决类似这样的问题，就可以利用 蓄水池算法(Reservoir Sampling)。

基本原理

假设需要采样的数量为 $ k $ 。

首先构建一个 $ k $ 个元素的数组，将序列的前 $ k $ 个元素放入数组中。
对于从第 $j $ 个元素 $(j>k)$ ，以 $\frac{k}{j}$ 的概率来决定该元素是否被替换到数组中，数组中的 $k$ 个元素被替换的概率是相同的。
当遍历完所有元素之后，数组中剩下的元素即为采样样本

证明

假设我们真的按照 $\frac{k}{j}$ 的概率来决定该元素是否被替换到数组中，有如下结论：

对于第 $i$ 个元素 $(i \le k)$ ：
- 在 $k$ 步之前，被选中的概率为 1。
- 当走到第 $k+1$ 步时
  
  第 i 个元素被第 k+1 个元素替换的概率 $=$ 第 k+1 个元素被选中的概率 $\times$ 第 i 个元素被选中替换的概率
  
  即:
  $$
  \frac{k}{k+1} \times \frac{1}{k}=\frac{1}{k+1}
  $$
  那么，第 $i$ 个元素不被第 $k+1$ 个元素替换的概率为 : $1-\frac{1}{k+1}=\frac{k}{k+1} $
  
  以此类推，第 $i$ 个元素不被第 $t$ ($t>k$)个元素替换掉的概率： $1-\frac{1}{t}=\frac{t-1}{t} $
- 当来到第 $n$ 个元素时：
  
  第 i 个元素被保留在水池中的概率 = 第 k+1 个元素被选中的概率 $\times$ 第 k+1 个元素不被替换的概率
  $$
  p_i = 1\times \frac{k}{k+1} \times \frac{k+1}{k+2} \times \frac{k+2}{k+3} \ldots \times \frac{n-1}{n}=\frac{k}{n}
  $$
对于第 $j$ 个元素 $(j>k)$：
- 在第 $j$ 步被选中的概率： $\frac{k}{j}$
- 第 $j+1$ 步没有被替换掉的概率： $1-\frac{k}{j+1} \times \frac{1}{k}=\frac{j}{j+1} $
- 当递推到第 $n$ 个元素时： 被保留的概率 = 被选中的概率 $\times$ 不被替换的概率，即:
  $$
  p_j= \frac{k}{j} \times \frac{j}{j+1} \times \frac{j+1}{j+2} \ldots \times \frac{n-1}{n}=\frac{k}{n}
  $$
因此对于每个元素，被保留的概率都是 $\frac{k}{n}$

应用示例

考虑长度为 $n$ 的数组，设定目标 $t$ ，要求便利数组过程中挑出和 $t$ 值相等的数字的下标，使得每个等于 $t$ 的值被选中的概率相等。

应用

遍历长度为n的数组。当第 $i$ 次遇到 $t$ 的元素时，随机选择区间 $[0, i)$ 内的一个整数，如果其等于0，则将返回值置为该元素的下标，否则返回值不变(也就是判定是否触发了 $\frac{1}{i}$ 的概率)

证明

设数组中有 $k$ 个值为 $t$ 的元素，该算法会保证这 $k$ 个元素的下标最终返回值概率均为 $1/k$ ，证明：

第i次遇到值为target的元素下标成为最终返回值的概率 = 第i次随机选择的值=0的概率 x 第 i+1 次随机选择的值!=0的概率x … x 第k次随机选择的值!=0的概率
$$
p_i= \frac{1}{i} \times\left(1-\frac{1}{i+1}\right) \times \ldots \times\left(1-\frac{1}{k}\right)=\frac{1}{i} \times \frac{i}{i+1} \times \ldots \times \frac{k-1}{k}=\frac{1}{k}
$$