每日算法 — 使用java实现蓄水池抽样：数据流环境下的无偏随机采样

引言

在大数据时代，数据往往以流的形式持续产生，总量可能达到 TB 甚至 PB 级别，无法全部加载到内存中处理。蓄水池抽样（Reservoir Sampling）是一类能够在仅遍历一次数据流的情况下，从中均匀随机抽取 k 个样本的经典算法。它的空间复杂度仅为 O(k)，与数据流总规模 N 无关，被广泛应用于数据库采样、日志分析、A/B 测试用户分配、流式推荐系统等场景。

本文将用 Java 实现蓄水池抽样的核心算法及其加权变种，深入讲解其数学原理、无偏性证明与实际应用。

核心问题定义

给定一个未知长度 N 的数据流（N 可能极大，甚至无限），从中等概率地抽取 k 个样本，要求：

1. 每个元素被选中的概率均为 k/N
2. 只能遍历数据一次（单遍扫描）
3. 内存使用量为 O(k)

算法原理

基础蓄水池抽样（k=1）

当只需要抽取 1 个样本时，算法极其简洁：

1. 保留第一个元素作为当前候选
2. 对于第 i 个元素（i ≥ 2），以 1/i 的概率替换当前候选
3. 遍历结束后，候选即为所求

无偏性证明：第 i 个元素最终被选中的概率 = 第 i 步被选中的概率 × 之后未被替换的概率

$$P(选中i) = \frac{1}{i} \times \prod_{j=i+1}^{N} \left(1 – \frac{1}{j}\right) = \frac{1}{i} \times \frac{i}{N} = \frac{1}{N}$$

扩展蓄水池抽样（k > 1）

1. 先将前 k 个元素放入蓄水池
2. 对于第 i 个元素（i > k），以 k/i 的概率决定是否将其放入蓄水池
3. 如果决定放入，随机替换蓄水池中的一个已有元素
4. 每个元素被选中的概率为 k/N

Java 实现

基础蓄水池抽样

import java.util.Random;

public class ReservoirSampling<T> {
    private final Random random;

    public ReservoirSampling() {
        this.random = new Random();
    }

    /**
     * 从数据流中抽取 k 个样本
     * @param stream 数据流迭代器
     * @param k 样本数量
     * @return 均匀随机样本数组
     */
    @SuppressWarnings("unchecked")
    public T[] sample(Iterable<T> stream, int k) {
        if (k <= 0) {
            throw new IllegalArgumentException("k must be positive");
        }

        // 蓄水池，存储选中的样本
        Object[] reservoir = new Object[k];
        int count = 0;

        // 第一步：填充蓄水池的前 k 个元素
        for (T item : stream) {
            if (count < k) {
                reservoir[count] = item;
                count++;
            } else {
                // 第二步：对第 count+1 个及以后的元素进行概率替换
                // 以 k/(count+1) 的概率决定是否保留当前元素
                int j = random.nextInt(count + 1);
                if (j < k) {
                    // 替换蓄水池中第 j 个元素
                    reservoir[j] = item;
                }
                count++;
            }
        }

        // 如果数据流长度小于 k，返回实际数量的样本
        if (count < k) {
            Object[] result = new Object[count];
            System.arraycopy(reservoir, 0, result, 0, count);
            return (T[]) result;
        }

        return (T[]) reservoir;
    }

    /**
     * 从数据流中抽取 1 个样本（k=1 的特化版本，更简洁）
     */
    public T sampleOne(Iterable<T> stream) {
        T result = null;
        int count = 0;

        for (T item : stream) {
            count++;
            // 以 1/count 的概率替换当前结果
            if (random.nextInt(count) == 0) {
                result = item;
            }
        }

        return result;
    }
}

加权蓄水池抽样

当数据流中每个元素有不同的权重，需要按权重比例抽样时，使用加权蓄水池抽样：

import java.util.Random;

public class WeightedReservoirSampling<T> {
    private final Random random;

    public WeightedReservoirSampling() {
        this.random = new Random();
    }

    /**
     * 加权蓄水池抽样：按权重比例抽取 k 个样本
     * 使用 A-ExpJ 算法（Efraimidis and Spirakis, 2006）
     * 每个元素的关键值为 u^(1/w)，其中 u 是 [0,1] 均匀随机数，w 是权重
     */
    public static class WeightedItem<T> {
        T item;
        double weight;

        public WeightedItem(T item, double weight) {
            if (weight <= 0) {
                throw new IllegalArgumentException("Weight must be positive");
            }
            this.item = item;
            this.weight = weight;
        }
    }

    @SuppressWarnings("unchecked")
    public T[] weightedSample(Iterable<WeightedItem<T>> stream, int k) {
        if (k <= 0) throw new IllegalArgumentException("k must be positive");

        // 使用优先队列维护蓄水池，按关键值排序
        // Java 的 PriorityQueue 默认是最小堆
        java.util.PriorityQueue<ScoreItem<T>> minHeap = 
            new java.util.PriorityQueue<>(k, (a, b) -> Double.compare(a.score, b.score));

        for (WeightedItem<T> weightedItem : stream) {
            // 生成关键值：u^(1/w)，u ~ Uniform(0,1)
            double u = random.nextDouble();
            double score = Math.pow(u, 1.0 / weightedItem.weight);

            if (minHeap.size() < k) {
                minHeap.offer(new ScoreItem<>(weightedItem.item, score));
            } else if (score > minHeap.peek().score) {
                // 当前元素的关键值大于蓄水池中最小的关键值，替换之
                minHeap.poll();
                minHeap.offer(new ScoreItem<>(weightedItem.item, score));
            }
        }

        T[] result = (T[]) new Object[minHeap.size()];
        int i = 0;
        for (ScoreItem<T> si : minHeap) {
            result[i++] = si.item;
        }
        return result;
    }

    private static class ScoreItem<T> {
        T item;
        double score;

        ScoreItem(T item, double score) {
            this.item = item;
            this.score = score;
        }
    }
}

分布式蓄水池抽样

在大数据分布式环境中，数据分散在多个节点上，需要支持合并多个本地样本：

import java.util.ArrayList;
import java.util.List;
import java.util.Random;

public class DistributedReservoirSampling<T> {
    private final Random random;

    public DistributedReservoirSampling() {
        this.random = new Random();
    }

    /**
     * 合并多个本地蓄水池样本
     * 每个本地样本都知道其对应的数据分片大小
     * @param localSamples 各节点的本地样本及对应数据量
     * @param k 最终样本数
     */
    @SuppressWarnings("unchecked")
    public T[] mergeSamples(List<LocalSample<T>> localSamples, int k) {
        // 计算总数据量
        long totalSize = localSamples.stream().mapToLong(ls -> ls.dataSize).sum();

        // 根据每个分片的大小比例，计算应从中抽取的样本数
        List<T> merged = new ArrayList<>();
        for (LocalSample<T> ls : localSamples) {
            int targetFromThis = (int) Math.round(k * (double) ls.dataSize / totalSize);
            // 从本地样本中再抽样
            for (int i = 0; i < Math.min(targetFromThis, ls.samples.length); i++) {
                merged.add(ls.samples[i]);
            }
        }

        // 如果合并后超过 k，进行截断
        if (merged.size() > k) {
            // 使用 Fisher-Yates shuffle 的前 k 个
            for (int i = merged.size() - 1; i > 0 && i >= merged.size() - k; i--) {
                int j = random.nextInt(i + 1);
                T temp = merged.get(i);
                merged.set(i, merged.get(j));
                merged.set(j, temp);
            }
            merged = merged.subList(merged.size() - k, merged.size());
        }

        return merged.toArray((T[]) new Object[0]);
    }

    public static class LocalSample<T> {
        T[] samples;
        long dataSize; // 该节点处理的数据总量

        public LocalSample(T[] samples, long dataSize) {
            this.samples = samples;
            this.dataSize = dataSize;
        }
    }
}

测试与验证

import java.util.*;

public class ReservoirSamplingDemo {
    public static void main(String[] args) {
        // 生成测试数据流：10000 个整数
        List<Integer> dataStream = new ArrayList<>();
        for (int i = 0; i < 10000; i++) {
            dataStream.add(i);
        }

        int k = 100;
        int iterations = 10000; // 重复实验次数
        int[] frequency = new int[10000];

        ReservoirSampling<Integer> sampler = new ReservoirSampling<>();

        System.out.println("=== 蓄水池抽样无偏性验证 ===\n");
        System.out.println("数据流大小: " + dataStream.size());
        System.out.println("样本数 k: " + k);
        System.out.println("重复实验次数: " + iterations);

        // 重复抽样，统计每个元素被选中的频率
        for (int iter = 0; iter < iterations; iter++) {
            // 每次打乱顺序模拟不同顺序的数据流
            Collections.shuffle(dataStream);
            Integer[] samples = sampler.sample(dataStream, k);
            for (Integer s : samples) {
                frequency[s]++;
            }
        }

        // 统计结果
        double expectedFrequency = iterations * k / 10000.0;
        System.out.println("\n期望每个元素被选中次数: " + String.format("%.2f", expectedFrequency));

        int minFreq = Integer.MAX_VALUE;
        int maxFreq = 0;
        double sumFreq = 0;
        for (int i = 0; i < 10000; i++) {
            minFreq = Math.min(minFreq, frequency[i]);
            maxFreq = Math.max(maxFreq, frequency[i]);
            sumFreq += frequency[i];
        }

        System.out.println("实际最小选中次数: " + minFreq);
        System.out.println("实际最大选中次数: " + maxFreq);
        System.out.println("实际平均选中次数: " + String.format("%.2f", sumFreq / 10000));
        System.out.println("偏差范围: " + String.format("%.2f%%", (maxFreq - minFreq) * 100.0 / expectedFrequency));

        // 加权抽样测试
        System.out.println("\n=== 加权蓄水池抽样测试 ===");
        List<WeightedReservoirSampling.WeightedItem<String>> weightedItems = Arrays.asList(
            new WeightedReservoirSampling.WeightedItem<>("A-高权重", 10.0),
            new WeightedReservoirSampling.WeightedItem<>("B-中权重", 5.0),
            new WeightedReservoirSampling.WeightedItem<>("C-低权重", 1.0)
        );

        WeightedReservoirSampling<String> weightedSampler = new WeightedReservoirSampling<>();
        Map<String, Integer> weightedCount = new HashMap<>();
        for (int i = 0; i < 10000; i++) {
            String[] result = weightedSampler.weightedSample(weightedItems, 1);
            for (String r : result) {
                weightedCount.merge(r.substring(0, 1), 1, Integer::sum);
            }
        }

        System.out.println("加权抽样结果分布:");
        weightedCount.forEach((key, count) -> 
            System.out.println("  " + key + ": " + count + " 次 (" + String.format("%.2f%%", count * 100.0 / 10000) + ")")
        );
        System.out.println("理论期望: A≈62.5%, B≈31.25%, C≈6.25%");
    }
}

复杂度分析

实际应用场景

1. 数据库采样：PostgreSQL 的 TABLESAMPLE SYSTEM 和 TABLESAMPLE BERNOULLI 底层使用蓄水池抽样变体
2. 日志分析：从海量日志流中实时抽取样本进行异常检测
3. A/B 测试：将在线用户按均匀概率分配到不同实验组
4. 推荐系统冷启动：从物品池中均匀采样初始化用户兴趣模型
5. 数据质量检测：对无法全量加载的大数据集进行抽样质检
6. 流式机器学习：在数据流中维护固定大小的训练样本窗口

算法变体

总结

蓄水池抽样的优美之处在于它以 O(k) 的空间复杂度解决了看似需要 O(N) 空间的问题。其核心思想——”以动态调整的概率维护一个固定大小的代表集合”——是流式算法设计的典范。掌握蓄水池抽样，意味着掌握了一种在资源受限条件下处理无限数据流的重要工具。