Java中灵活高效实现加权概率分布的通用方法

Random random = new Random();
int randomInt = random.nextInt(10) + 1; // 生成1到10的随机数

if (randomInt <= 3) { // 概率 3/10
    System.out.println("Magnificent!");
} else if (randomInt >= 7) { // 概率 4/10 (7,8,9,10)
    System.out.println("Marvelous!");
} else { // 概率 3/10 (4,5,6)
    System.out.println("Delectable!");
}

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 static class WeightedValue<T> {
        final double weight; // 权重
        final T value;      // 对应的值

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

    // 比较器，用于按权重降序排序 WeightedValue 对象
    private final Comparator<WeightedValue<T>> byWeight =
            Comparator.comparing((WeightedValue<T> wv) -> wv.weight).reversed(); // 降序

    // 使用 TreeSet 存储 WeightedValue，自动按权重降序排序
    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("WeightedRandom set is empty.");
        }

        // 生成一个介于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()); // 如果随机数大于或等于当前累积权重，则继续

        return result.value;
    }
}

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

        // 添加带权重的值
        // "AAA" 权重为3 (相当于30%)
        // "BBB" 权重为2 (相当于20%)
        // "CCC" 权重为5 (相当于50%)
        randomSelector.put(3, "AAA");
        randomSelector.put(2, "BBB");
        randomSelector.put(5, "CCC");

        // 模拟1000次选择，观察分布情况
        System.out.println("--- 模拟1000次加权随机选择 ---");
        int countAAA = 0;
        int countBBB = 0;
        int countCCC = 0;

        for (int i = 0; i < 1000; i++) {
            String value = randomSelector.next();
            // System.out.println(value); // 可以取消注释查看每次选择结果

            switch (value) {
                case "AAA":
                    countAAA++;
                    break;
                case "BBB":
                    countBBB++;
                    break;
                case "CCC":
                    countCCC++;
                    break;
            }
        }

        System.out.println("选择结果统计 (1000次):");
        System.out.printf("AAA: %d 次 (%.2f%%)%n", countAAA, (double) countAAA / 1000 * 100);
        System.out.printf("BBB: %d 次 (%.2f%%)%n", countBBB, (double) countBBB / 1000 * 100);
        System.out.printf("CCC: %d 次 (%.2f%%)%n", countCCC, (double) countCCC / 1000 * 100);

        // 示例：移除或添加新的权重，验证灵活性
        System.out.println("\n--- 移除并重新添加权重后 ---");
        // 这里只是演示，实际WeightedRandom类没有提供移除方法，
        // 如果需要动态调整，可能需要清空并重新put，或者扩展WeightedRandom类
        // 假设我们重新创建一个实例来模拟权重变化
        WeightedRandom<String> dynamicSelector = new WeightedRandom<>();
        dynamicSelector.put(1, "Alpha");
        dynamicSelector.put(9, "Beta"); // Beta现在有更高的权重

        System.out.println("一次动态选择结果: " + dynamicSelector.next());
    }
}