LeetCode [528] Random Pick with Weight

528. Random Pick with Weight

Medium

Given an array w of positive integers, where w[i] describes the weight of index i(0-indexed), write a function pickIndex which randomly picks an index in proportion to its weight.

For example, given an input list of values w = [2, 8], when we pick up a number out of it, the chance is that 8 times out of 10 we should pick the number 1 as the answer since it's the second element of the array (w[1] = 8).

 

Example 1:

Input
["Solution","pickIndex"]
[[[1]],[]]
Output
[null,0]

Explanation
Solution solution = new Solution([1]);
solution.pickIndex(); // return 0. Since there is only one single element on the array the only option is to return the first element.

Example 2:

Input
["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
[[[1,3]],[],[],[],[],[]]
Output
[null,1,1,1,1,0]

Explanation
Solution solution = new Solution([1, 3]);
solution.pickIndex(); // return 1. It's returning the second element (index = 1) that has probability of 3/4.
solution.pickIndex(); // return 1
solution.pickIndex(); // return 1
solution.pickIndex(); // return 1
solution.pickIndex(); // return 0. It's returning the first element (index = 0) that has probability of 1/4.

Since this is a randomization problem, multiple answers are allowed so the following outputs can be considered correct :
[null,1,1,1,1,0]
[null,1,1,1,1,1]
[null,1,1,1,0,0]
[null,1,1,1,0,1]
[null,1,0,1,0,0]
......
and so on.

 

Constraints:

• 1 <= w.length <= 10000
• 1 <= w[i] <= 10^5
• pickIndex will be called at most 10000 times.

class Solution {
    Random random;
    int[] wSum;
    int n, L;
    public Solution(int[] w) {
        random = new Random();
        wSum = w;
        n = w.length;
        for(int i=1; i<n; ++i){
            wSum[i] += wSum[i-1];
        }
        L = wSum[n-1];
    }
    
    /*
    [2,5,3,4]
     0     1   2    3  -- m
    12 34567 890 1234  -- wSum
    */
    public int pickIndex() {
        int rnd = random.nextInt(L) + 1;//nextInt return [0, L-1]
        //find index m such that
        //1. wSum[m]>=rnd
        //2. wSum[m-1]<rnd
        int l = 0, r = n-1;
        while(l<=r){
            int m = (l+r)/2;
            if(wSum[m]<rnd){
                l = m+1;
            }else{//wSum[m]>=rnd
                if(m==0) return m;
                else if(wSum[m-1]<rnd) return m;
                else{//wSum[m-1]>=rnd
                    r = m-1;
                }
            }
        }
        return -1;
    }
}

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

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18

class Solution { TreeMap<Integer, Integer> map = new TreeMap<>(); int s; Random rand = new Random(); public Solution(int[] w) { s = 0; for(int i=0; i<w.length; ++i){ s += w[i]; map.put(s, i); } } public int pickIndex() { int r = rand.nextInt(s)+1; Map.Entry<Integer, Integer> t = map.ceilingEntry(r); return t.getValue(); } }