Problem
Given an array w
of positive integers, where w[i]
describes the weight of index i
, write a function pickIndex
which randomly picks an index in proportion to its weight.
Note:
1 <= w.length <= 10000
1 <= w[i] <= 10^5
pickIndex
will be called at most10000
times.
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]
Explanation of Input Syntax:
The input is two lists: the subroutines called and their arguments. Solution
‘s constructor has one argument, the array w
. pickIndex
has no arguments. Arguments are always wrapped with a list, even if there aren’t any.
Solution: Binary Search
- Convert PDF to CDF
- Uniformly sample a value s in [1, sum(weights)].
- Use binary search to find first index such that PDF[index] >= s.
Time complexity: Init O(n), query O(logn)
Space complexity: O(n)
C++
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
// Author: Huahua // Running time: 184 ms class Solution { public: Solution(vector<int> w) { sums_ = std::move(w); for (int i = 1; i < sums_.size(); ++i) sums_[i] += sums_[i - 1]; } int pickIndex() { int s = rand() / static_cast<double>(RAND_MAX) * sums_.back() + 1; return lower_bound(sums_.begin(), sums_.end(), s) - sums_.begin(); // or // int l = 0; // int r = sums_.size(); // while (l < r) { // int m = (r - l) / 2 + l; // if (sums_[m] < s) // l = m + 1; // else // r = m; // } // return l; } private: vector<int> sums_; }; |