按权重随机选择

题目链接

按权重随机选择

题目描述

注意点

• 1 <= w.length <= 10^4
• 1 <= w[i] <= 10^5

解答思路

• 因为本题选取i的概率取决与i对应的元素值，假设所以元素之和为total，可以按顺序将每个元素按其权值分布在total的相应区间上，在随机选择时，选取[1, total]之间的随机数，在二分查找找到该随机数处于哪个区间即可
• 为了方便根据元素权值分布到total上的某段区间，采用前缀和，例如：[3,1,2,4]，前缀和分别为[3, 4, 6, 10]，第一个元素分布在[1, 3]，第二个元素分布在[4, 4]，第三个元素分布在[5, 6]，第四个元素分布在[7, 10]，同一区间内不会有两个及以上元素，每个元素按照权值分配相应空间

代码

class Solution {
    // 从第一个元素开始，维护了一个区间代表该范围内取该整数，区间大小就等于该元素的权值
    // 前缀和，preSum[i]表示前i个整数的和
    int[] preSum;
    int total;
    Random random;

    public Solution(int[] w) {
        preSum = new int[w.length];
        preSum[0] = w[0];
        for (int i = 1; i < w.length; i++) {
            preSum[i] = preSum[i - 1] + w[i];
        }
        total = preSum[w.length - 1];
        random = new Random();
    }
    
    public int pickIndex() {
        int x = random.nextInt(total) + 1;
        return binarySearch(x);
    }

    public int binarySearch(int x) {
        int l = 0;
        int r = preSum.length - 1;
        while (l <= r) {
            int mid = l + ((r - l) >> 1);
            if (preSum[mid] == x) {
                return mid;
            }
            if (preSum[mid] < x) {
                l = mid + 1;
            } else {
                r = mid - 1;
            }
        }
        return l;
    }
}

/**
 * Your Solution object will be instantiated and called as such:
 * Solution obj = new Solution(w);
 * int param_1 = obj.pickIndex();
 */

关键点

• 前缀和将元素按照权值分布到相应区间上
• 二分查找的思想