528. Random Pick with Weight

1. Description

You are given a 0-indexed array of positive integers w where w[i] describes the weight of the ith index.
You need to implement the function pickIndex(), which randomly picks an index in the range [0, w.length - 1] (inclusive) and returns it. The probability of picking an index i is w[i] / sum(w).
For example, if w = [1, 3], the probability of picking index 0 is 1 / (1 + 3) = 0.25 (i.e., 25%), and the probability of picking index 1 is 3 / (1 + 3) = 0.75 (i.e., 75%).

2. Example

Example 1:
Input
[“Solution”,“pickIndex”]
[[[1]],[]]
Output
[null,0]

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 is returning the second element (index = 1) that has a probability of 3/4. solution.pickIndex(); // return 1
solution.pickIndex(); // return 1
solution.pickIndex(); // return 1
solution.pickIndex(); // return 0. It is returning the first element (index = 0) that has a probability of 1/4.

Since this is a randomization problem, multiple answers are allowed.
All of 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.

3. Constraints

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

4. Solutions

Prefix Sum & Binary Search

n = weights.size()

class Solution {
private:
    vector<int> weight_prefix_sum;

public:
    // Solution(vector<int> &w)
    Solution(vector<int> &weights) {
        // Time complexity: O(n)  
        // Space complexity: O(n)
        weight_prefix_sum = vector<int>(weights.size());
        weight_prefix_sum[0] = weights[0];
        for (int i = 1; i < weights.size(); ++i) {
            weight_prefix_sum[i] = weights[i] + weight_prefix_sum[i - 1];
        }

        srand(time(0));
    }

    int pickIndex() {
        // Time complexity: O(logn)  
        // Space complexity: O(1)
        int random_number = random();

        cout << random_number << endl;

        int result = 0;
        for (int left = 0, right = weight_prefix_sum.size(); left < right;) {
            int mid = (left + right) / 2;

            random_number <= weight_prefix_sum[mid] ? result = mid, right = mid : left = mid + 1;
        }

        return result;
    }
};