程式語言 - LeetCode - C++ - 528. Random Pick with Weight



題目:


解答:

class Solution {
private:
    int total;
    vector<int> prefix;

public:
    Solution(vector<int>& w) {
        total = 0;

        for (int n : w) {
            total += n;
            prefix.push_back(total);
        }
    }
    
    int pickIndex() {
        int l = 0;
        int r = prefix.size() - 1;
        int target = rand() % total + 1;

        while (l < r) {
            int m = l + ((r - l) >> 1);

            if (prefix[m] < target) {
                l = m +1;
            }
            else {
                r = m;
            }
        }

        return l;
    }
};

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