Random Pick with Weight

Problem

Random Pick with Weight

Approach

Different numbers have different weight. The weight also can be saw as a small range in a big range. So I will generate some number randomly and then check the number will be included in which range.

Suppose the w array is [3, 2], then the random number obtained from the array can be {0, 1, 2, 3, 4, 5}. And more specifically, numbers of {0,1,2} stand for the element of index 0, which is 3 in w; numbers of {3,4} stand for the element of index 1 in array w, which is 2.

Code

class Solution {
    
    private int[] w;
    private Random rand;
    
    public Solution(int[] w) {
        this.w = new int[w.length];
        for(int i = 0; i < w.length; i ++){
            this.w[i] = w[i];
            if(i >= 1){
                this.w[i] += this.w[i - 1];
            }
        }
        rand = new Random();
    }
    
    private int binarySearch(int target){
        int l = 0, r = w.length - 1;
        while(l <= r){
            int mid = l + (r - l) / 2;
            if(w[mid] < target){
                l = mid + 1;
            } else {
                r = mid - 1;
            }
        }
        if(w[l] == target){
            return l + 1;
        }
        return l;
    }
    
    public int pickIndex() {
        int rId = rand.nextInt(w[w.length - 1]);
        return binarySearch(rId);
    }
}

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