蓄水池抽样算法（Reservoir Sampling）
许多年以后，当听说蓄水池抽样算法时，邱simple将会想起，那个小学数学老师带他做“小明对水池边加水边放水，求何时能加满水”应用题的下午。
一、问题
我是在一次失败的面试经历中听说蓄水池算法的。之后上网搜了搜，知道是一个数据抽样算法，寥寥几行，却暗藏玄机。主要用来解决如下问题。
给定一个数据流，数据流长度N很大，且N直到处理完所有数据之前都不可知，请问如何在只遍历一遍数据（O(N)）的情况下，能够随机选取出m个不重复的数据。
这个场景强调了3件事：
- 数据流长度N很大且不可知，所以不能一次性存入内存。
- 时间复杂度为O(N)。
- 随机选取m个数，每个数被选中的概率为m/N。
第1点限制了不能直接取N内的m个随机数，然后按索引取出数据。第2点限制了不能先遍历一遍，然后分块存储数据，再随机选取。第3点是数据选取绝对随机的保证。讲真，在不知道蓄水池算法前，我想破脑袋也不知道该题做何解。
二、核心代码及原理
蓄水池抽样算法的核心如下：
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。
算法思路大致如下：
- 如果接收的数据量小于m，则依次放入蓄水池。
- 当接收到第i个数据时，i >= m，在[0, i]范围内取以随机数d，若d的落在[0, m-1]范围内，则用接收到的第i个数据替换蓄水池中的第d个数据。
- 重复步骤2。
算法的精妙之处在于：当处理完所有的数据时，蓄水池中的每个数据都是以m/N的概率获得的。
下面用白话文推导验证该算法。假设数据开始编号为1.
第i个接收到的数据最后能够留在蓄水池中的概率=第i个数据进入过蓄水池的概率*之后第i个数据不被替换的概率（第i+1到第N次处理数据都不会被替换）。
- 当i<=m时，数据直接放进蓄水池，所以第i个数据进入过蓄水池的概率=1。
- 当i>m时，在[1,i]内选取随机数d，如果d<=m，则使用第i个数据替换蓄水池中第d个数据，因此第i个数据进入过蓄水池的概率=m/i。
- 当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。
- 当i>m时，程序从接收到第i+1个数据时开始有可能替换第i个数据。则参考上述第3点，之后第i个数据不被替换的概率=i/N。
- 结合第1点和第3点可知，当i<=m时，第i个接收到的数据最后留在蓄水池中的概率=1*m/N=m/N。结合第2点和第4点可知，当i>m时，第i个接收到的数据留在蓄水池中的概率=m/i*i/N=m/N。综上可知，每个数据最后被选中留在蓄水池中的概率为m/N。
这个算法建立在统计学基础上，很巧妙地获得了“m/N”这个概率。
三、深入一些——分布式蓄水池抽样（Distributed/Parallel Reservoir Sampling）
一块CPU的计算能力再强，也总有内存和磁盘IO拖他的后腿。因此为提高数据吞吐量，分布式的硬件搭配软件是现在的主流。
如果遇到超大的数据量，即使是O(N)的时间复杂度，蓄水池抽样程序完成抽样任务也将耗时很久。因此分布式的蓄水池抽样算法应运而生。运作原理如下：
- 假设有K台机器，将大数据集分成K个数据流，每台机器使用单机版蓄水池抽样处理一个数据流，抽样m个数据，并最后记录处理的数据量为N1, N2, ..., Nk, ..., NK(假设m
- 取[1, N]一个随机数d，若d
- 取[1, N]一个随机数d，若d
m/N的概率验证如下：
- 第k台机器中的蓄水池数据被选取的概率为m/Nk。
- 从第k台机器的蓄水池中选取一个数据放进最终蓄水池的概率为Nk/N。
- 第k台机器蓄水池的一个数据被选中的概率为1/m。（不放回选取时等概率的）
- 重复m次选取，则每个数据被选中的概率为m*(m/Nk*Nk/N*1/m)=m/N。
四、算法验证
写一份完整的代码，用来验证蓄水池抽样的随机性。数据集大小为1000，蓄水池容量为10，做10_0000次抽样。如果程序正确，那么每个数被抽中的次数接近1000次。
package cn.edu.njupt.qyz;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.LinkedList;
import java.util.List;
import java.util.Random;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
public class ReservoirSampling {
static ExecutorService exec = Executors.newFixedThreadPool(4);
// 抽样任务，用作模拟并行抽样
private static class SampleTask implements Callable
{
// 输入该任务的数据
private int[] innerData;
// 蓄水池容量
private int m;
SampleTask (int m, int[] innerData)
{
this.innerData = innerData;
this.m = m;
}
@Override
public int[] call() throws Exception
{
int[] reservoir = sample(this.m, this.innerData);
return reservoir;
}
}
// 并行抽样
public static int[] mutiSample(int m, int[] dataStream) throws InterruptedException, ExecutionException
{
Random rand = new Random();
int[] reservoir = initReservoir(m, dataStream);
// 生成3个范围内随机数，将数据切成4份
List list = getRandInt(rand, dataStream.length);
int s1 = list.get(0);
int s2 = list.get(1);
int s3 = list.get(2);
// 每个任务处理的数据量
double n1 = s1 - 0;
double n2 = s2 - s1;
double n3 = s3 - s2;
double n4 = dataStream.length - s3;
// 并行抽样
Future f1 = exec.submit(new SampleTask(m, Arrays.copyOfRange(dataStream, 0, s1)));
Future f2 = exec.submit(new SampleTask(m, Arrays.copyOfRange(dataStream, s1, s2)));
Future f3 = exec.submit(new SampleTask(m, Arrays.copyOfRange(dataStream, s2, s3)));
Future f4 = exec.submit(new SampleTask(m, Arrays.copyOfRange(dataStream, s3, dataStream.length)));
List r1 = getList(f1.get());
List r2 = getList(f2.get());
List r3 = getList(f3.get());
List r4 = getList(f4.get());
// 进行m次抽样
for (int i = 0; i < m; i++)
{
int p = rand.nextInt(dataStream.length);
// 根据随机数落在的范围选择元素
if (p < s1)
{
reservoir[i] = getRandEle(r1, rand.nextInt(r1.size()));
}
else if (p < s2)
{
reservoir[i] = getRandEle(r2, rand.nextInt(r2.size()));
}
else if (p < s3)
{
reservoir[i] = getRandEle(r3, rand.nextInt(r3.size()));
}
else
{
reservoir[i] = getRandEle(r4, rand.nextInt(r4.size()));
}
}
return reservoir;
}
// 根据输入返回随机位置的元素，并且删除该元素，模拟不放回
private static int getRandEle(List list, int idx)
{
return list.remove(idx);
}
// 获取bound范围内的3个随机数，用来分割数据集
private static List getRandInt(Random rand, int bound)
{
Set set = new TreeSet<>();
List list = new ArrayList<>();
while (set.size() < 3)
{
set.add(rand.nextInt(bound));
}
for (int e: set)
{
list.add(e);
}
return list;
}
// 数据转换成List
private static List getList(int[] arr)
{
List list = new LinkedList<>();
for (int a : arr)
{
list.add(a);
}
return list;
}
// 单机版蓄水池抽样
public static int[] sample(int m, int[] dataStream)
{
// 随机数生成器，以系统当前nano时间作为种子
Random rand = new Random();
int[] reservoir = initReservoir(m, dataStream);
// 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];
}
}
return reservoir;
}
private static int[] initReservoir (int m, int[] dataStream)
{
int[] reservoir;
if (m > dataStream.length)
{
reservoir = new int[dataStream.length];
}
else
{
reservoir = new int[m];
}
return reservoir;
}
// 单机版测试
public void test()
{
// 样本长度
int len = 1000;
// 蓄水池容量
int m = 10;
// 抽样次数，用作验证抽样的随机性
int iterTime = 100000;
// 每个数字被抽到的次数
int[] freq = new int[len];
// 样本
int[] dataStream = new int[len];
// init dataStream
for (int i = 0; i < dataStream.length; i++)
{
dataStream[i] = i;
}
// count freq
for (int k = 0; k < iterTime; k++)
{
// 进行抽样
int[] reservoir = sample(m, dataStream);
// 计算出现次数
for (int i = 0; i < reservoir.length; i++)
{
int ele = reservoir[i];
freq[ele] += 1;
}
}
printStaticInfo(freq);
}
// 测试并行抽样
public void mutiTest() throws InterruptedException, ExecutionException
{
// 样本长度
int len = 1000;
// 蓄水池容量
int m = 10;
// 抽样次数，用作验证抽样的随机性
int iterTime = 10_0000;
// 每个数字被抽到的次数
int[] freq = new int[len];
// 样本
int[] dataStream = new int[len];
// init dataStream
for (int i = 0; i < dataStream.length; i++)
{
dataStream[i] = i;
}
// count freq
for (int k = 0; k < iterTime; k++)
{
// 进行抽样
int[] reservoir = mutiSample(m, dataStream);
// 计算出现次数
for (int i = 0; i < reservoir.length; i++)
{
int ele = reservoir[i];
freq[ele] += 1;
}
}
printStaticInfo(freq);
}
// 打印统计信息
private void printStaticInfo (int[] freq)
{
// 期望、方差和标准差
double avg = 0;
double var = 0;
double sigma = 0;
// print
for (int i = 0; i < freq.length; i++)
{
if (i % 10 == 9) System.out.println();
System.out.print(freq[i] + ", ");
avg += ((double)(freq[i]) / freq.length);
var += (double)(freq[i] * freq[i]) / freq.length;
}
// 输出统计信息
System.out.println("\n===============================");
var = var - avg * avg;
sigma = Math.sqrt(var);
System.out.println("Average: " + avg);
System.out.println("Variance: " + var);
System.out.println("Standard deviation: " + sigma);
}
public static void main (String[] args) throws InterruptedException, ExecutionException
{
ReservoirSampling rs = new ReservoirSampling();
rs.mutiTest();
}
}
单机版输出和并行版的输出类似，截取片段如下：
948, 1006, 1014, 1019, 1033, 1040, 948, 1014, 1000, 951,
1014, 987, 1049, 1043, 1034, 983, 1006, 974, 1060, 1009,
986, 1021, 1024, 963, 1041, 1028, 988, 1011, 975, 980,
1055, 1017, 1010, 1018, 1013, 983, 942, 1056, 1003, 1063,
1004, 1004, 999, 976, 957, 935, 1061, 1018, 1002, 1018,
1019, 946, 985, 1057, 1012, 965, 978, 1040, 1026, 1064,
1026, 1018, 980, 996, 1025, 1028, 1006, 944, 986, 981,
923, 1015, 991, 1019, 1024, 1143, 989, 985, 1022, 1019,
1004, 1000, 989, 972, 1041, 988, 1050, 932, 975, 1037,
1016, 983, 1051, 1003, 983, 986, 1017, 1009, 936, 993,
965, 976, 1001, 1000, 988, 1030, 1050, 1024, 981, 985,
935, 1023, 996, 1007, 1013, 1046, 1003, 1006, 973, 989,
943,
===============================
Average: 1000.0000000000002
Variance: 1011.8799999983748
Standard deviation: 31.81006130139291
此外，为了对比单机版与并行版（4线程）的性能差异，使用10_0000大小的数据集，蓄水池容量10，进行100_0000次重复抽样，对比两者的运行时间。结果如下
---------单机版----------
===============================
Average: 100.00000000000125
Variance: 100.31497999751264
Standard deviation: 10.015736617818613
---------并行版----------
===============================
Average: 100.00000000000169
Variance: 100.63045999737915
Standard deviation: 10.031473470900432
单机版耗时：2006s
并行版耗时：1265s
从输出结果可以看出，算法保证了数据选取的随机性。且并行版算法能够有效提高数据吞吐量。
五、应用场景
蓄水池抽样的O(N)时间复杂度，O(m)空间复杂度令其适用于对流数据、大数据集的等概率抽样。比如一个大文本数据，随机输出其中的几行。
六、总结
象征性总结：优雅巧妙的算法——蓄水池抽样。
七、参考文献
- 数据工程师必知算法：蓄水池抽样
- 【算法34】蓄水池抽样算法 (Reservoir Sampling Algorithm)
- 分布式/并行蓄水池抽样 (Distributed/Parallel Reservoir Sampling)
- Distributed/Parallel Reservoir Sampling