Java中实现灵活且简洁的加权随机选择

import java.util.Comparator;
import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.ThreadLocalRandom;

public class WeightedRandom<T> {    
    // 用于按权重降序排序的比较器
    private final Comparator<WeightedValue<T>> byWeight = 
        Comparator.comparing((WeightedValue<T> wv) -> wv.weight).reversed();

    // 存储加权值的TreeSet，自动按权重降序排列
    private final Set<WeightedValue<T>> weightedValues = 
        new TreeSet<>(byWeight);

    // 所有权重的总和
    private double totalWeight;

    /**
     * 添加一个加权值到集合中。
     * 权重必须大于0。
     *
     * @param weight 值的权重
     * @param value  要添加的值
     */
    public void put(double weight, T value) {
        if (weight <= 0) {
            // 忽略非正权重，或者可以抛出异常
            return;
        }
        totalWeight += weight;
        weightedValues.add(new WeightedValue<>(weight, value));
    }

    /**
     * 根据权重随机选择一个值。
     *
     * @return 随机选中的值
     * @throws NoSuchElementException 如果没有添加任何加权值
     */
    public T next() {
        if (weightedValues.isEmpty()) {
            throw new NoSuchElementException("No weighted values have been added.");
        }

        // 生成一个介于0到totalWeight之间的随机数
        double rnd = ThreadLocalRandom.current().nextDouble(totalWeight);
        double sum = 0; // 累积权重

        Iterator<WeightedValue<T>> iterator = weightedValues.iterator();
        WeightedValue<T> result;

        // 遍历加权值，直到累积权重超过随机数
        do {
            result = iterator.next();
            sum += result.weight;
        } while (rnd >= sum && iterator.hasNext()); // 注意这里是rnd >= sum，确保在边界情况下也能正确选择

        return result.value;
    }

    /**
     * 内部静态类，用于封装值及其权重。
     */
    private static class WeightedValue<T> {
        final double weight;
        final T value;

        public WeightedValue(double weight, T value) {
            this.weight = weight;
            this.value = value;
        }
    }
}

public class WeightedRandomExample {
    public static void main(String[] args) {
        WeightedRandom<String> randomSelector = new WeightedRandom<>();

        // 添加加权值：
        // "AAA" 权重为 3
        // "BBB" 权重为 2
        // "CCC" 权重为 5
        randomSelector.put(3, "AAA"); // 3 / (3+2+5) = 30% 概率
        randomSelector.put(2, "BBB"); // 2 / (3+2+5) = 20% 概率
        randomSelector.put(5, "CCC"); // 5 / (3+2+5) = 50% 概率

        System.out.println("进行1000次加权随机选择：");
        // 模拟1000次随机选择，观察结果分布
        int countA = 0;
        int countB = 0;
        int countC = 0;

        for (int i = 0; i < 1000; i++) {
            String value = randomSelector.next();
            // System.out.println(value); // 如果需要打印每次结果
            switch (value) {
                case "AAA":
                    countA++;
                    break;
                case "BBB":
                    countB++;
                    break;
                case "CCC":
                    countC++;
                    break;
            }
        }

        System.out.println("AAA 出现次数: " + countA + " (期望: 300)");
        System.out.println("BBB 出现次数: " + countB + " (期望: 200)");
        System.out.println("CCC 出现次数: " + countC + " (期望: 500)");
        System.out.println("总次数: " + (countA + countB + countC));
    }
}