Description

Submission

class Solution {
    vector<int> presum;
    int n;
    std::default_random_engine engine;
    std::uniform_int_distribution<int> uniform_dist;
public:
    Solution(vector<int>& w) {
        n = w.size();
        presum.resize(n);
        presum[0] = w[0];
        for(int i = 1; i < n; ++i) {
            presum[i] = presum[i-1] + w[i];
        }

        random_device r;
        engine = default_random_engine(r());
        uniform_dist = std::uniform_int_distribution<int>(1, presum[n-1]);
    }
    
    int pickIndex() {
        int idx = uniform_dist(engine);
        // cout << idx << endl;
        auto it = lower_bound(presum.begin(), presum.end(), idx);

        return it-presum.begin();
        // return 0;
    }
};

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