Given an array w of positive integers, where w[i] describes the weight of index i, write a functionpickIndex which randomly picks an index in proportion to its weight.
Note:
1 <= w.length <= 100001 <= w[i] <= 10^5pickIndex will be called at most 10000 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.
[code lang="java"]
class Solution {
public Solution(int[] w) {
}
public int pickIndex() {
}
}
/**
* Your Solution object will be instantiated and called as such:
* Solution obj = new Solution(w);
* int param_1 = obj.pickIndex();
*/
[code]
Idea – 1
Consider the sequence

, here at indices 0, 1, 2, and 3 we have weights 1, 3, 2, and 4. Thus 10% of the time we must pick index 0, 30% of the time we must pick index 1, 20% of the time we must pick index 2, and 40% of the time we must pick index 3. For our example totalWeight = 1+3+2+4 = 10. We divide the total weight like: [1, 1], [2, 4], [5, 6], and [7, 10]. Now if we draw x from [1..10] at random, we pick 0 if x is in [1, 1], we pick 1 if x is in [2, 4] etc. So we pick an index proportionally to its weight. Time complexity is

, space complexity is

.
[code lang="java"]
class Solution {
private static final Random rng = new Random(System.currentTimeMillis()%Integer.MAX_VALUE);
private int totalWeight = 0;
private int[] weights;
public Solution(int[] w) {
for(int x : w)
{
totalWeight += x;
}
weights = w;
}
public int pickIndex() {
int x = 1+rng.nextInt(totalWeight);
int cumsum = 0;
for(int i = 0; i < weights.length; ++i)
{
cumsum += weights[i];
if(x <= cumsum)
{
return i;
}
}
return -1;
}
}
[code]
Runtime: 90 ms, faster than 13.88% of Java online submissions for Random Pick with Weight.Memory Usage: 52 MB, less than 35.80% of Java online submissions for Random Pick with Weight.
Idea – 2
If we precompute the cumsum, we can dispense with the summing in pickIndex. Time complexity remains same, space complexity becomes

.
[code lang="java"]
class Solution {
private static final Random rng = new Random(System.currentTimeMillis()%Integer.MAX_VALUE);
private int[] cumsum;
public Solution(int[] w) {
cumsum = new int[w.length];
cumsum[0] = w[0];
for(int i = 1; i < w.length; ++i)
{
cumsum[i] = cumsum[i-1] + w[i];
}
}
public int pickIndex() {
int x = 1+rng.nextInt(cumsum[cumsum.length-1]);
for(int i = 0; i < cumsum.length; ++i)
{
if(x <= cumsum[i])
{
return i;
}
}
return -1;
}
}
[code]
Runtime: 80 ms, faster than 27.95% of Java online submissions for Random Pick with Weight.Memory Usage: 52.5 MB, less than 34.07% of Java online submissions for Random Pick with Weight.
Idea – 3
Since cumsum is sorted, we could leverage binary search in pickIndex which reduces time complexity to

, space complexity remains

.
[code lang="java"]
class Solution {
private static final Random rng = new Random(System.currentTimeMillis()%Integer.MAX_VALUE);
private int[] cumsum;
public Solution(int[] w) {
cumsum = new int[w.length];
cumsum[0] = w[0];
for(int i = 1; i < w.length; ++i)
{
cumsum[i] = cumsum[i-1] + w[i];
}
}
public int pickIndex() {
int x = 1+rng.nextInt(cumsum[cumsum.length-1]);
return ceilSearch(cumsum, x);
}
private int ceilSearch(int[] A, int key)
{
int lo = 0, hi = A.length-1;
while(lo <= hi)
{
int mid = lo+(hi-lo)/2;
if(key == A[mid])
{
return mid;
}
else if(key < A[mid])
{
hi = mid-1;
}
else
{
lo = mid+1;
}
}
return lo;
}
}
[code]
Runtime: 67 ms, faster than 92.00% of Java online submissions for Random Pick with Weight.Memory Usage: 44.3 MB, less than 95.91% of Java online submissions for Random Pick with Weight.