528. Random Pick with Weight

Solution

class Solution {
    int[] prefixSum;
    int totalSum;
    public Solution(int[] w) {
        prefixSum = new int[w.length];
        int curSum = 0;
        for (int i = 0; i < w.length; i++) {
            curSum += w[i];
            prefixSum[i] = curSum;
        }
        totalSum = curSum;
    }
    
    public int pickIndex() {
        double pickedIdx = Math.random() * totalSum; // Math.random generates [0,1)

        // in the prefixsum array, find the first element > pickedIdx
        int left = 0, right = prefixSum.length - 1;
        while (left <= right) {
            int mid = left + (right - left) / 2;
            if (prefixSum[mid] == pickedIdx) {
                left = mid + 1;
            } else if (prefixSum[mid] > pickedIdx) {
                right = mid - 1;
            } else {
                left = mid + 1;
            }
        }
        return left;
    }
}