如何仅依据各类元素的数量实现高效无放回抽样？

朴素实现的核心瓶颈 #

现有方案需要把所有带重复次数的元素全量展开成列表再打乱，时间和空间复杂度完全和所有元素的总权重值正相关：如果单个元素的出现次数是100万，就要在内存中生成100万条重复的元素引用，遍历构造列表的耗时也会随总权重线性增长，在权重量级较大的场景下效率极低。

高适配性的优化方案 #

最常用的改进思路是无需展开全量元素，基于动态权重的逐次随机选择逻辑完成抽样，抽样概率分布和全量展开的朴素方案完全等价，不存在偏差。常见实现有两类：

1. 基础版：前缀和遍历实现 #

核心流程：

• 初始化时直接保留原始的elems数组和对应的剩余数量sizes数组，计算当前所有剩余数量的总和totalWeight
• 每一轮抽样执行固定操作：
  • 生成一个取值范围为[0, totalWeight)的随机整数randNum
  • 遍历所有元素，逐个累加对应元素的剩余数量currentSize，当累加值首次大于randNum时，当前遍历到的元素即为本次抽中结果
  • 抽中后将该元素对应的剩余数量减1，同时将totalWeight减1
• 重复上述流程k次，即可得到k次无放回抽样的结果

该方案实现逻辑简单，空间复杂度为O(n)（n为不同元素的种类数），完全和单个元素的权重大小无关，哪怕单个元素权重是1000万，也不会增加额外内存开销。

对应的优化后Java代码示例：

import java.util.Arrays;
import java.util.Random;

public class WeightedSampling {
    public static void main(String[] args) {
        String[] elems = { "A", "B", "C", "D", "E" };
        Integer[] sizes = { 10, 5, 4, 7, 3 };
        int k = 3;
        // 记录当前总剩余权重
        int totalWeight = Arrays.stream(sizes).mapToInt(Integer::intValue).sum();
        Random random = new Random();

        for (int sampleIdx = 0; sampleIdx < k; sampleIdx++) {
            // 生成[0, totalWeight)范围内的随机数
            int randNum = random.nextInt(totalWeight);
            int accumWeight = 0;
            String selectedElem = "";
            // 遍历累加权重匹配随机数
            for (int i = 0; i < elems.length; i++) {
                accumWeight += sizes[i];
                if (accumWeight > randNum) {
                    selectedElem = elems[i];
                    // 抽中后减少对应元素的剩余数量
                    sizes[i]--;
                    break;
                }
            }
            System.out.println(selectedElem);
            totalWeight--;
        }
    }
}

2. 进阶版：树状数组/线段树优化 #

如果不同元素的种类数n特别大（比如上万个类别），逐次遍历累加权重的复杂度会上升到O(kn)，此时可以用树状数组（二叉索引树）* 存储每个元素的剩余数量：

• 初始化时把每个元素的剩余数量写入树状数组，整体初始化复杂度为O(n logn)
• 每次查询前缀和匹配随机数、更新对应位置的剩余数量时，都可以通过二分查找+树状数组的前缀和查询能力，把单轮抽样的复杂度降到O(logn)
• 整体抽样流程的总复杂度为O(n logn + k logn)，和元素权重大小完全解耦，适合超大规模类别、高权重的抽样场景。

方案对比 #

内容的提问来源于stack exchange，提问作者molfo