You are given a 0-indexed array of positive integers w
where w[i]
describes the weight of the i<sup>th</sup>
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 index0
is1 / (1 + 3) = 0.25
(i.e.,25%
), and the probability of picking index1
is3 / (1 + 3) = 0.75
(i.e.,75%
).
Test Cases
Example 1:
**Input**
["Solution","pickIndex"]
[[[1]],[]]
**Output**
[null,0]
**Explanation**
Solution solution = new Solution([1]);
solution.pickIndex(); // return 0. The only option is to return 0 since there is only one element in w.
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.
Constraints:
1 <= w.length <= 10<sup>4</sup>
1 <= w[i] <= 10<sup>5</sup>
pickIndex
will be called at most10<sup>4</sup>
times.
Solution
class Solution {
private int[] prefix;
private Random random;
public Solution(int[] w) {
random = new Random();
prefix = new int[w.length];
prefix[0] = w[0];
for(int i=1; i<w.length; i++) {
prefix[i] = w[i] + prefix[i-1];
}
}
public int pickIndex() {
int limit = prefix[prefix.length - 1];
int digit = random.nextInt(limit) + 1;
return search(digit);
}
private int search(int digit) {
int left = 0, right = prefix.length - 1;
while(left <= right) {
int mid = left + (right - left)/2;
if (prefix[mid] == digit) {
return mid;
}
if (prefix[mid] < digit) {
left = mid+1;
} else {
right = mid-1;
}
}
return left;
}
}
/**
* Your Solution object will be instantiated and called as such:
* Solution obj = new Solution(w);
* int param_1 = obj.pickIndex();
*/