528. Random Pick with Weight

class Solution {
    private final int[] prefixSumArray;

    public Solution(int[] w) {
        this.prefixSumArray = new int[w.length];
        prefixSumArray[0] = w[0];
        for (int i = 1; i < prefixSumArray.length; i++) {
            prefixSumArray[i] = prefixSumArray[i - 1] + w[i];
        }
    }

    public int pickIndex() {
        // nums: [4, 2, 9]
        // prefixSumArray: [4, 6, 15]
        int randomValue = ThreadLocalRandom.current().nextInt(prefixSumArray[prefixSumArray.length - 1]) + 1; // [1, prefixSumArray[prefixSumArray.length - 1]]

        int left = 0, right = prefixSumArray.length - 1;
        while (left < right) {
            int mid = (left + right) >>> 1;
            if (prefixSumArray[mid] < randomValue) {
                left = mid + 1; // 可以排除当前区间，所以 left = mid + 1
            } else {
                right = mid; // 只能排除当前区间右侧的区间，所以 right = mid
            }
        }

        return left;
    }
}

注意前缀和数组中的值为各个权重区间的右边界值，如：nums = [1, 3, 1] 时，prefixSumArray = [1, 4, 5]。

References

528. Random Pick with Weight
剑指 Offer II 071. 按权重生成随机数