如何根据指定概率分布从数组中随机选取对应元素

实现思路 #

核心是用经典的带权重随机抽样算法，逻辑非常简单：

• 先定义和目标数组长度一致的权重数组，每个位置的权重对应你需要的概率百分比，总和为100即可
• 计算权重的前缀和数组，相当于把0-100的数值区间切割成和数组元素一一对应的子区间，每个子区间的长度就是对应元素的抽取概率
• 生成一个0到100之间的随机浮点数，判断该数值落在哪个前缀区间，返回对应索引的元素即可

代码实现（Java版本，匹配你给出的数组语法） #

import java.util.Random;

public class WeightedRandom {
    // 目标数组
    private static final int[] goals = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
    // 硬编码权重数组，对应各元素的抽取概率百分比，总和为100
    private static final double[] weights = {50, 30, 12, 2.5, 0.85, 0.3, 0.2, 0.1, 0.03, 0.01, 0.01};
    // 前缀和缓存数组
    private static final double[] prefixSum;
    private static final Random random = new Random();

    // 初始化前缀和数组，程序启动只算一次
    static {
        prefixSum = new double[weights.length];
        prefixSum[0] = weights[0];
        for (int i = 1; i < weights.length; i++) {
            prefixSum[i] = prefixSum[i-1] + weights[i];
        }
    }

    public static int pick() {
        // 生成0~100之间的随机浮点数
        double randomVal = random.nextDouble() * 100;
        // 匹配对应的区间
        for (int i = 0; i < prefixSum.length; i++) {
            if (randomVal < prefixSum[i]) {
                return goals[i];
            }
        }
        // 极端情况兜底，正常不会走到这里
        return goals[goals.length - 1];
    }

    // 测试用方法，运行可以看到实际抽样概率和设定值吻合
    public static void main(String[] args) {
        int[] count = new int[goals.length];
        int testTimes = 100000;
        for (int i = 0; i < testTimes; i++) {
            count[pick()]++;
        }
        for (int i = 0; i < goals.length; i++) {
            System.out.printf("元素%d 实际抽样概率：%.2f%%%n", goals[i], count[i] * 100.0 / testTimes);
        }
    }
}

补充说明 #

上面代码里weights数组的后几位是我临时补的数值凑满100%，如果你有固定的概率递推规则，直接修改weights数组的对应值即可，不需要调整其他逻辑。如果后续数组长度变大，可以把遍历匹配区间的逻辑换成二分查找，时间复杂度可以从O(n)降到O(logn)，当前11个元素的场景用遍历就足够高效。

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