528 Random Pick with Weight

528. Random Pick with Weight

1. Question

Given an arraywof positive integers, wherew[i]describes the weight of indexi, write a functionpickIndex which randomly picks an index in proportion to its weight.

Note:

1. 1 <= w.length <= 10000

2. 1 <= w[i] <= 10^5

3. pickIndexwill be called at most10000times.

Example 1:

Input: 
["Solution","pickIndex"]
[[[1]],[]]

Output: 
[null,0]

Example 2:

Input: 
["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
[[[1,3]],[],[],[],[],[]]

Output: 
[null,0,1,1,1,0]

2. Implementation

(1) TreeMap

class Solution {
    TreeMap<Integer, Integer> map;
    Random rand;
    int totalWeight;
    int[] weight;

    public Solution(int[] w) {
        this.weight = w;
        map = new TreeMap();
        totalWeight = 0;
        rand = new Random();

        for (int i = 0; i < weight.length; i++) {
            totalWeight += weight[i];
            map.put(totalWeight, i);
        }
    }

    public int pickIndex() {
        int randWeight = rand.nextInt(totalWeight) + 1;
        int key = map.ceilingKey(randWeight);
        return map.get(key);
    }
}

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

(2) Binary Search

class Solution {
    Random rand;
    int[] weight;
    int totalWeight;

    public Solution(int[] w) {
        rand = new Random();
        this.weight = w;
        totalWeight = weight[0];

        for (int i = 1; i < weight.length; i++) {
            totalWeight += weight[i];
            weight[i] += weight[i - 1];
        }
    }

    public int pickIndex() {
        int target = rand.nextInt(totalWeight) + 1;

        int left = 0, right = weight.length - 1;

        while (left + 1 < right) {
            int mid = left + (right - left) / 2;

            if (weight[mid] < target) {
                left = mid + 1;
            }
            else {
                right = mid;
            }
        }
        return weight[left] >= target ? left : right;
    }
}

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

3. Time & Space Complexity

(1) TreeMap: Time: constructor: O(n), pickIndex: O(logn), Space: O(n)

(2) Binary Tree: Time: constructor: O(n), pickIndex: O(logn), Space: O(n)