每日算法 — 使用java实现蓄水池抽样：数据流随机采样与无偏估计

蓄水池抽样（Reservoir Sampling）是Vitter于1985年提出的一类随机抽样算法，它解决了在数据流或超大规模数据集中进行均匀随机抽样的核心难题：当数据总量N未知或极大（无法一次载入内存）时，如何以O(1)空间复杂度获取k个均匀分布的无偏样本？蓄水池抽样的精妙之处在于，它只需遍历数据一次，且每个元素被抽中的概率严格等于k/N，数学上完全无偏。

核心算法与数学证明

基本蓄水池抽样（k=1）：维护一个候选样本，遍历到第i个元素时，以1/i的概率将其替换为当前元素。最终每个元素被选中的概率均为1/N。

归纳证明：设第i个元素最终保留为样本的概率为P(i)。它必须在第i轮被选中（概率1/i），且在后续每轮j>i中不被替换。在第j轮不被替换的概率为 $(j-1)/j$。因此：

$$P(i) = \frac{1}{i} \times \frac{i}{i+1} \times \frac{i+1}{i+2} \times \cdots \times \frac{N-1}{N} = \frac{1}{N}$$

扩展蓄水池抽样（k>1）：维护一个大小为k的蓄水池。遍历到第i个元素时（i>k），以k/i的概率决定是否将其纳入蓄水池，并随机替换蓄水池中的一个现有样本。

Java完整实现

import java.util.*;

/**
 * 蓄水池抽样实现 - 数据流随机采样算法
 * 在数据总量未知或极大的场景下，以O(k)空间获取k个均匀无偏样本
 */
public class ReservoirSampling {

    private final Random random;

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

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

    /**
     * 基础蓄水池抽样：从数据流中抽取k个均匀随机样本
     * 时间复杂度：O(N)，空间复杂度：O(k)
     * 每个元素被选中的概率严格等于 k/N
     *
     * @param stream 数据流迭代器
     * @param k 样本容量
     * @param <T> 元素类型
     * @return 包含k个均匀随机样本的列表
     */
    public <T> List<T> sample(Iterator<T> stream, int k) {
        if (k <= 0) return new ArrayList<>();

        List<T> reservoir = new ArrayList<>(k);
        int count = 0;

        // 第一阶段：填充蓄水池（前k个元素直接放入）
        while (count < k && stream.hasNext()) {
            reservoir.add(stream.next());
            count++;
        }

        // 若数据总量不足k，直接返回所有数据
        if (count < k) {
            return reservoir;
        }

        // 第二阶段：遍历剩余元素，以概率k/i决定是否替换
        while (stream.hasNext()) {
            T item = stream.next();
            count++;
            // 生成[0, count)范围内的随机整数
            int j = random.nextInt(count);
            // 若j < k，则用当前元素替换蓄水池中第j个元素
            if (j < k) {
                reservoir.set(j, item);
            }
        }

        return reservoir;
    }

    /**
     * 从List中抽取k个随机样本的便捷方法
     */
    public <T> List<T> sample(List<T> data, int k) {
        return sample(data.iterator(), k);
    }

    /**
     * 加权蓄水池抽样：每个元素带有权重，按权重比例抽样
     * 采用A-Res算法：维护一个优先级队列，元素的优先级为 u^(1/w)
     * 其中u为[0,1]均匀随机数，w为元素权重
     *
     * @param stream 带权重的数据流（Map.Entry<元素, 权重>）
     * @param k 样本容量
     * @param <T> 元素类型
     * @return 按权重比例抽取的k个样本
     */
    public <T> List<T> weightedSample(Iterator<Map.Entry<T, Double>> stream, int k) {
        if (k <= 0) return new ArrayList<>();

        // 使用最小堆维护k个最高优先级的元素
        PriorityQueue<WeightedItem<T>> heap = new PriorityQueue<>(
            Comparator.comparingDouble(item -> item.priority)
        );

        while (stream.hasNext()) {
            Map.Entry<T, Double> entry = stream.next();
            T item = entry.getKey();
            double weight = entry.getValue();

            if (weight <= 0) continue;

            // 计算优先级：u^(1/weight)，u ~ Uniform(0,1)
            double u = random.nextDouble();
            double priority = Math.pow(u, 1.0 / weight);

            if (heap.size() < k) {
                heap.offer(new WeightedItem<>(item, priority));
            } else if (priority > heap.peek().priority) {
                // 当前元素优先级高于堆中最小元素，替换之
                heap.poll();
                heap.offer(new WeightedItem<>(item, priority));
            }
        }

        List<T> result = new ArrayList<>();
        for (WeightedItem<T> item : heap) {
            result.add(item.value);
        }
        return result;
    }

    /**
     * 加权抽样内部辅助类
     */
    private static class WeightedItem<T> {
        T value;
        double priority;

        WeightedItem(T value, double priority) {
            this.value = value;
            this.priority = priority;
        }
    }

    /**
     * 分布式蓄水池抽样合并：将多个子蓄水池合并为全局蓄水池
     * 适用于MapReduce等分布式计算场景
     *
     * @param subReservoirs 各分区抽取的子样本列表
     * @param k 全局样本容量
     * @param <T> 元素类型
     * @return 合并后的全局均匀样本
     */
    @SafeVarargs
    public final <T> List<T> mergeReservoirs(int k, List<T>... subReservoirs) {
        // 收集所有子蓄水池的元素，并记录其来源的分区大小
        List<T> allItems = new ArrayList<>();
        for (List<T> sub : subReservoirs) {
            allItems.addAll(sub);
        }

        // 若子样本总量不足k，直接返回（实际场景中通常k <= 子样本总量）
        if (allItems.size() <= k) {
            return allItems;
        }

        // 从合并后的元素中重新执行蓄水池抽样
        Collections.shuffle(allItems, random);
        return allItems.subList(0, k);
    }

    /**
     * 验证抽样的均匀性：统计大量重复实验中的选择频率
     * 理论上每个元素被选中的次数应接近 (k/N) * trials
     */
    public void uniformityTest(int n, int k, int trials) {
        int[] selectionCount = new int[n];

        for (int t = 0; t < trials; t++) {
            List<Integer> data = new ArrayList<>();
            for (int i = 0; i < n; i++) data.add(i);
            Collections.shuffle(data);

            List<Integer> sample = sample(data, k);
            for (int idx : sample) {
                selectionCount[idx]++;
            }
        }

        double expected = (double) k * trials / n;
        System.out.println("均匀性验证 (N=" + n + ", k=" + k + ", trials=" + trials + ")");
        System.out.println("理论期望选中次数: " + String.format("%.2f", expected));
        System.out.println("实际选中次数分布:");

        double sumSqDiff = 0;
        for (int i = 0; i < n; i++) {
            double diff = selectionCount[i] - expected;
            sumSqDiff += diff * diff;
            if (i < 5 || i >= n - 5) {
                System.out.println("  元素 " + i + ": " + selectionCount[i] +
                    " (偏差 " + String.format("+%.2f%%", (selectionCount[i] - expected) / expected * 100) + ")");
            }
        }
        double rmse = Math.sqrt(sumSqDiff / n);
        System.out.println("RMSE: " + String.format("%.2f", rmse) +
            " (RMSE/Expected=" + String.format("%.4f", rmse / expected) + ")");
    }

    /**
     * 主方法：演示蓄水池抽样的核心功能
     */
    public static void main(String[] args) {
        ReservoirSampling sampler = new ReservoirSampling(42);

        // 场景1：从大数据流中抽取样本
        System.out.println("=== 场景1：从100万条数据流中抽取10个样本 ===");
        List<Integer> largeStream = new ArrayList<>();
        for (int i = 0; i < 1_000_000; i++) {
            largeStream.add(i);
        }
        Collections.shuffle(largeStream);

        List<Integer> sample = sampler.sample(largeStream, 10);
        System.out.println("抽取样本: " + sample);

        // 场景2：加权抽样
        System.out.println("\n=== 场景2：按权重比例抽样 ===");
        List<Map.Entry<String, Double>> weightedData = new ArrayList<>();
        weightedData.add(new AbstractMap.SimpleEntry<>("A", 10.0));
        weightedData.add(new AbstractMap.SimpleEntry<>("B", 5.0));
        weightedData.add(new AbstractMap.SimpleEntry<>("C", 3.0));
        weightedData.add(new AbstractMap.SimpleEntry<>("D", 2.0));

        Map<String, Integer> freq = new HashMap<>();
        for (int i = 0; i < 10000; i++) {
            List<String> ws = sampler.weightedSample(weightedData.iterator(), 1);
            for (String s : ws) {
                freq.merge(s, 1, Integer::sum);
            }
        }
        System.out.println("10000次加权抽样频率:");
        for (Map.Entry<String, Integer> e : freq.entrySet()) {
            System.out.println("  " + e.getKey() + ": " + e.getValue() +
                " (" + String.format("%.2f%%", e.getValue() / 100.0) + ")");
        }

        // 场景3：均匀性验证
        System.out.println("\n=== 场景3：均匀性统计验证 ===");
        sampler.uniformityTest(100, 5, 10000);
    }
}

复杂度分析

与先打乱再取前k个的Fisher-Yates洗牌法相比，蓄水池抽样不需要O(N)的额外空间存储索引或数据，是数据流场景的唯一可行方案。

算法扩展与变体

工业级应用

• 日志采样：大型分布式系统从海量日志流中实时抽取样本用于监控和告警
• A/B测试：从活跃用户流中实时随机划分实验组和对照组
• 数据库查询优化：Oracle、PostgreSQL等使用蓄水池抽样实现TABLESAMPLE子句
• 流式机器学习：从无限数据流中维护固定大小的训练样本窗口
• 网络流量分析：从高速网络包流中抽取样本进行深度检测