528. Random Pick with Weight

You are given an array of positive integers w where w[i] describes the weight of ith index (0-indexed).

We need to call the function pickIndex() which randomly returns an integer in the range [0, w.length - 1]. pickIndex() should return the integer proportional to its weight in the w array. For example, for w = [1, 3], the probability of picking the index 0 is 1 / (1 + 3) = 0.25 (i.e 25%) while the probability of picking the index 1 is 3 / (1 + 3) = 0.75 (i.e 75%).

More formally, the probability of picking index i is w[i] / sum(w).

class Solution {
    map<int, int> m;
    int tot = 0;
public:
    Solution(vector<int>& w) {
        for(int i = 0; i < w.size(); i++) {
            m[tot] = i;
            tot += w[i];
        }
    }

    int pickIndex() {
        int num = rand() % tot;
        return prev(m.upper_bound(num))->second;
    }
};

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