在大数据时代，数据往往以流的形式持续产生：服务器日志实时写入、用户行为事件不断上报、传感器数据持续采集。一个常见需求是从总长度未知的数据流中，均匀随机地抽取固定数量的样本。若先存储全部数据再抽样，内存开销巨大甚至不可行。蓄水池抽样（Reservoir Sampling）正是解决这一问题的经典随机化算法——它仅用O(k)的额外空间，即可在单次遍历中完成公平采样，并严格保证每个元素被选中的概率完全相等。本文将用Java完整实现基础蓄水池抽样、加权扩展版本以及分布式场景下的合并策略，并给出严格的概率证明。
一、问题建模与核心思想
1.1 为什么需要蓄水池抽样
假设某直播平台需要在一场没有预先知道总时长的活动中，从所有实时弹幕中随机抽取10名幸运观众发放奖品。问题的难点在于：
- 数据总量N未知：活动结束前不知道会有多少条弹幕；
- 只能顺序访问：数据像水流一样逐个到来，无法回退；
- 内存受限：不可能先把所有弹幕存下来再抽样；
- 公平性要求：每条弹幕对应的观众被选中的概率必须严格相等。
朴素思路”先存后抽”在空间上不可行，而”随机决定是否保留每个元素”又难以保证最终恰好选出k个样本。蓄水池抽样的核心洞察是：维护一个容量为k的”蓄水池”，当第i个元素到来时，以k/i的概率决定是否将其放入池中，并随机替换池中的一个旧元素。这样巧妙的替换规则，恰好能保证每个元素最终留在池中的概率都是k/N。
1.2 算法形式化描述
输入：一个长度未知（或极大）的数据流 stream，样本容量 k。
输出：从中均匀随机选出的 k 个元素。
Algorithm R（基础蓄水池抽样）：
1. 将数据流的前 k 个元素直接放入蓄水池；
2. 对于第 i 个元素（i 从 k+1 开始计数）：
– 生成一个 [1, i] 范围内的随机整数 j；
– 如果 j <= k，则用第 i 个元素替换蓄水池中第 j 个位置的元素；
– 否则不做任何操作；
3. 数据流结束后，蓄水池中的 k 个元素即为所求样本。
二、基础蓄水池抽样的Java实现
2.1 泛型蓄水池类设计
为了通用性，我们使用Java泛型实现一个可复用的蓄水池抽样器。核心数据结构是一个大小为 k 的数组（蓄水池），以及一个计数器记录已处理的数据总量。
import java.util.Random;
/**
* 基础蓄水池抽样器（Algorithm R）
* 适用于从总长度未知的数据流中均匀随机抽取k个样本
* @param <T> 数据流元素类型
*/
public class ReservoirSampler<T> {
// 蓄水池容量
private final int k;
// 蓄水池数组
private final Object[] reservoir;
// 当前已处理的数据总量
private long count;
// 随机数生成器
private final Random random;
/**
* 构造指定容量的蓄水池
* @param k 需要抽取的样本数量，必须大于0
*/
public ReservoirSampler(int k) {
if (k <= 0) {
throw new IllegalArgumentException("蓄水池容量k必须大于0");
}
this.k = k;
this.reservoir = new Object[k];
this.count = 0;
this.random = new Random();
}
/**
* 处理流中的单个元素
* 前k个元素直接放入蓄水池，后续元素以k/count概率决定是否替换
* @param item 数据流中的当前元素
*/
public void feed(T item) {
count++;
if (count <= k) {
// 阶段一：蓄水池未满，直接放入
reservoir[(int) (count - 1)] = item;
} else {
// 阶段二：蓄水池已满，以 k/count 的概率决定是否替换
// 生成 [0, count-1] 的随机数，若落在 [0, k-1] 范围内则替换
int j = random.nextInt((int) count);
if (j < k) {
reservoir[j] = item;
}
}
}
/**
* 批量处理数据流
* @param stream 可迭代的数据源
*/
public void feedAll(Iterable<T> stream) {
for (T item : stream) {
feed(item);
}
}
/**
* 获取当前蓄水池中的样本
* @return 长度为k的样本数组（可能包含null如果输入不足k个元素）
*/
@SuppressWarnings("unchecked")
public T[] getSample() {
T[] result = (T[]) new Object[k];
System.arraycopy(reservoir, 0, result, 0, k);
return result;
}
/**
* 获取已处理的数据总量
*/
public long getTotalCount() {
return count;
}
}
2.2 概率正确性验证程序
为了直观验证算法的公平性，我们构造一个测试：从1000个编号元素中抽取10个样本，重复10000次实验，统计每个元素被选中的频次。理想情况下，每个元素的选中次数应接近 10000 * 10 / 1000 = 100 次。
import java.util.Arrays;
/**
* 蓄水池抽样正确性验证
* 通过大量重复实验，验证每个元素被选中的概率是否均匀
*/
public class ReservoirValidation {
public static void main(String[] args) {
final int N = 1000; // 数据流总长度
final int k = 10; // 样本容量
final int trials = 10000; // 重复实验次数
int[] frequency = new int[N]; // 记录每个元素被选中的总次数
for (int t = 0; t < trials; t++) {
ReservoirSampler<Integer> sampler = new ReservoirSampler<>(k);
// 模拟数据流：0, 1, 2, ..., 999
for (int i = 0; i < N; i++) {
sampler.feed(i);
}
Integer[] sample = sampler.getSample();
for (Integer val : sample) {
if (val != null) {
frequency[val]++;
}
}
}
// 统计分析
double expected = (double) trials * k / N;
double sumSqDiff = 0;
int minFreq = Integer.MAX_VALUE;
int maxFreq = Integer.MIN_VALUE;
for (int freq : frequency) {
sumSqDiff += Math.pow(freq - expected, 2);
minFreq = Math.min(minFreq, freq);
maxFreq = Math.max(maxFreq, freq);
}
double variance = sumSqDiff / N;
double stdDev = Math.sqrt(variance);
System.out.println("=== 蓄水池抽样均匀性验证 ===");
System.out.printf("数据流长度 N=%d, 样本容量 k=%d, 实验次数=%d%n", N, k, trials);
System.out.printf("理论期望选中次数: %.2f%n", expected);
System.out.printf("实际最小选中次数: %d%n", minFreq);
System.out.printf("实际最大选中次数: %d%n", maxFreq);
System.out.printf("标准差: %.4f (越小越均匀)%n", stdDev);
System.out.printf("前10个元素的选中频次: %s%n",
Arrays.toString(Arrays.copyOf(frequency, 10)));
}
}
三、概率证明：为什么每个元素被选中的概率恰好是 k/N
蓄水池抽样的优美之处在于其简洁的代码背后隐藏着严谨的概率平衡。我们用数学归纳法证明：对于任意第 m 个元素（1 <= m <= N），它在算法结束后留在蓄水池中的概率恰好是 k/N。
基础情况：当 m <= k 时，第 m 个元素一开始就被放入蓄水池。它在后续过程中被替换出去的概率是各次替换概率的并集：
$$P(\text{被替换}) = \sum_{i=k+1}^{N} P(\text{第i轮选中替换}) \cdot P(\text{恰好替换到m})$$
第 i 轮（i > k）以 k/i 的概率决定替换，且替换到第 m 个位置的概率是 1/k，因此第 m 个元素在第 i 轮被替换的概率是 1/i。它在所有轮次中存活下来的概率为：
$$P(\text{第m个元素最终保留}) = \prod_{i=k+1}^{N} \left(1 – \frac{1}{i}\right) = \frac{k}{N}$$
归纳步骤：当 m > k 时，第 m 个元素首先要被选中进入蓄水池（概率 k/m），然后在之后的 N - m 轮中不被替换出去。对于任意后续轮次 i > m，它在该轮被替换的概率是 1/i，因此存活概率为：
$$P(\text{第m个元素最终保留}) = \frac{k}{m} \cdot \prod_{i=m+1}^{N} \left(1 – \frac{1}{i}\right) = \frac{k}{m} \cdot \frac{m}{N} = \frac{k}{N}$$
由此得证：无论元素出现在数据流的哪个位置，其最终被选中的概率严格等于 k/N，满足均匀随机抽样的要求。
四、加权蓄水池抽样扩展
在实际应用中，不同元素往往具有不同的重要性权重。例如，在抽奖活动中VIP用户的权重可能更高。加权蓄水池抽样要求每个元素 e_i 被选中的概率正比于其权重 w_i。
4.1 A-Res算法核心思想
我们采用经典的A-Res算法（Efraimidis and Spirakis, 2006）。其核心洞察是：为每个元素计算一个关键字 key = u^(1/w)，其中 u 是 (0,1) 区间均匀随机数，w 是元素权重。然后维护一个容量为 k 的最小堆，保留关键字最大的 k 个元素。数学上可以证明，这种方案恰好能实现按权重比例抽样。
import java.util.PriorityQueue;
/**
* 加权蓄水池抽样器（A-Res算法）
* 每个元素附带权重，被选中的概率正比于权重
* @param <T> 元素类型
*/
public class WeightedReservoirSampler<T> {
/**
* 带权重的元素封装
*/
static class WeightedItem<T> {
final T item;
final double weight;
// 计算关键字：u^(1/weight)，u为均匀随机数
final double key;
WeightedItem(T item, double weight, Random random) {
if (weight <= 0) {
throw new IllegalArgumentException("权重必须为正数");
}
this.item = item;
this.weight = weight;
double u = random.nextDouble();
// 避免u=0导致key=0的边界情况
while (u == 0.0) u = random.nextDouble();
this.key = Math.pow(u, 1.0 / weight);
}
}
private final int k;
private final PriorityQueue<WeightedItem<T>> minHeap;
private final Random random;
private long count;
public WeightedReservoirSampler(int k) {
this.k = k;
// 按key升序排列的最小堆，堆顶是k个元素中key最小的
this.minHeap = new PriorityQueue<>((a, b) -> Double.compare(a.key, b.key));
this.random = new Random();
this.count = 0;
}
/**
* 处理带权重的元素
* @param item 数据元素
* @param weight 该元素的权重（必须 > 0）
*/
public void feed(T item, double weight) {
count++;
WeightedItem<T> weighted = new WeightedItem<>(item, weight, random);
if (minHeap.size() < k) {
// 蓄水池未满，直接加入
minHeap.offer(weighted);
} else if (weighted.key > minHeap.peek().key) {
// 当前元素的关键字大于堆顶，替换之
minHeap.poll();
minHeap.offer(weighted);
}
// 否则当前元素被丢弃
}
/**
* 获取最终的加权样本
*/
@SuppressWarnings("unchecked")
public T[] getSample() {
T[] result = (T[]) new Object[minHeap.size()];
int idx = 0;
for (WeightedItem<T> wi : minHeap) {
result[idx++] = wi.item;
}
return result;
}
}
五、分布式蓄水池抽样与合并策略
在分布式系统（如MapReduce、Flink、Spark Streaming）中，数据流被分割到多个节点并行处理。每个节点可以独立运行蓄水池抽样，但如何合并多个节点的局部样本得到全局均匀样本？
5.1 加权合并算法
假设有 m 个节点，节点 j 处理了 n_j 条数据并产出了容量为 k 的局部样本。合并时，我们可以将每个局部样本视为一个”元元素”，其权重为该节点处理的数据量 n_j，然后运行一次加权蓄水池抽样：
/**
* 分布式蓄水池抽样合并器
* 将多个节点的局部样本合并为全局均匀样本
*/
public class DistributedReservoirMerger<T> {
private final WeightedReservoirSampler<T> merger;
public DistributedReservoirMerger(int k) {
this.merger = new WeightedReservoirSampler<>(k);
}
/**
* 合并一个节点产生的局部样本
* @param localSample 该节点的k个局部样本
* @param nodeCount 该节点处理的数据总量
*/
public void mergeNodeSample(T[] localSample, long nodeCount) {
for (T item : localSample) {
if (item != null) {
// 以节点处理的数据量作为权重
merger.feed(item, (double) nodeCount);
}
}
}
public T[] getGlobalSample() {
return merger.getSample();
}
}
该合并策略的正确性基于一个关键观察：每个局部样本中的元素，其代表的全局选中概率与该节点处理的数据量成正比。通过以 n_j 为权重的二次抽样，最终恰好还原了全局均匀分布。
六、应用场景：幸运抽奖系统实战
我们将上述算法封装为一个完整的抽奖系统，模拟一场10万人在线的直播中抽取20名幸运观众。
import java.util.ArrayList;
import java.util.List;
/**
* 直播抽奖系统演示
* 模拟从大量实时用户中公平抽取获奖者
*/
public class LotterySystem {
static class User {
final String userId;
final String nickname;
final int level; // VIP等级，用于加权抽奖
User(String userId, String nickname, int level) {
this.userId = userId;
this.nickname = nickname;
this.level = level;
}
@Override
public String toString() {
return String.format("%s(%s, VIP%d)", nickname, userId, level);
}
}
public static void main(String[] args) {
final int TOTAL_USERS = 100_000;
final int PRIZE_COUNT = 20;
// 生成模拟用户流
List<User> userStream = new ArrayList<>();
for (int i = 0; i < TOTAL_USERS; i++) {
int vipLevel = (i % 100 == 0) ? 3 : (i % 10 == 0) ? 1 : 0;
userStream.add(new User("U" + i, "用户" + i, vipLevel));
}
System.out.println("=== 基础公平抽奖（人人平等） ===");
ReservoirSampler<User> fairSampler = new ReservoirSampler<>(PRIZE_COUNT);
fairSampler.feedAll(userStream);
User[] fairWinners = fairSampler.getSample();
for (User u : fairWinners) {
System.out.println(" 获奖者: " + u);
}
System.out.println("\n=== VIP加权抽奖（VIP权重 = 1 + level * 2） ===");
WeightedReservoirSampler<User> vipSampler = new WeightedReservoirSampler<>(PRIZE_COUNT);
for (User u : userStream) {
double weight = 1.0 + u.level * 2.0;
vipSampler.feed(u, weight);
}
User[] vipWinners = vipSampler.getSample();
for (User u : vipWinners) {
System.out.println(" 获奖者: " + u);
}
}
}
七、复杂度分析与算法对比
| 指标 | 基础蓄水池抽样 | 加权蓄水池抽样 | 先存后抽 |
|---|---|---|---|
| 时间复杂度 | O(N) | O(N log k) | O(N) |
| 空间复杂度 | O(k) | O(k) | O(N) |
| 遍历次数 | 1次 | 1次 | 2次 |
| 是否支持流式 | 是 | 是 | 否 |
| 是否支持加权 | 否 | 是 | 是（需二次加权） |
| 随机数需求 | N-k 个 | N 个 | k 个 |
八、总结
本文系统讲解了蓄水池抽样算法及其Java实现，核心内容包括：
- Algorithm R：用O(k)空间在单次遍历中完成均匀随机抽样，通过k/i的替换概率实现完美公平性；
- 概率证明：利用数学归纳法严格证明了每个元素最终保留概率等于k/N；
- A-Res加权扩展：通过
u^(1/w)关键字和最小堆，实现按权重比例的随机抽样； - 分布式合并：以各节点处理数据量为权重，二次抽样合并得到全局均匀样本；
- 工程实践：封装为可复用的泛型类，并应用于直播抽奖系统的完整示例。
蓄水池抽样的价值不仅在于其简洁的实现，更在于它揭示了随机化算法处理大规模数据时的核心哲学——用概率的精确性换取空间的极致压缩。读者可在此基础上继续探索：带删除操作的动态蓄ervoir抽样、滑动窗口上的蓄水池抽样、以及结合布隆过滤器进行去重后的联合抽样等高级场景。
