[Uber][LeetCode]Random Pick with Weight

Problem: You are given a 0-indexed array of positive integers w where w[i] describes the weight of the ith 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 index 0 is 1 / (1 + 3) = 0.25 (i.e., 25%), and the probability of picking index 1 is 3 / (1 + 3) = 0.75 (i.e., 75%).

Example:

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.

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.

Approach: One way to solve the problem is to create the array of length equals to the sum of all the numbers in the array w. Now put every index 'i' in this new array w[i] times and then just pick the index randomly in this new array and return the index. Let's take the above example [1, 3], the new array will be:

newArr = [0, 1, 1, 1]

So you can see now the if we can just pick an index randomly say 'randomIndex' and return the newArr[randomIndex], we have solved the problem and with good time complexity. The problem with this solution is the memory. Sum of all the numbers could be very large and taking an array of this size can be a problem so what we can do next?

We can use prefix sum. We can have a prefix sum array but with this we can't return a random index in constant time, instead we can use binary search to get the target in the prefix sum array. For above example our prefix sum array will be:

prefixSumArr = [1, 4]

Now we can generate a random number in between 1 to 4 and try to get its most suitable index in this sorted array using binary search. That's all!

Implementation in C#:

public class Solution

{

public Solution(int[] w)

    {

this.arr = new int[w.Length];

this.arr[0] = w[0];

for (int i = 1; i < w.Length; ++i)

{

this.arr[i] = this.arr[i - 1] + w[i];

}

this.rnd = new Random();

}

public int PickIndex()

    {

int target = rnd.Next(this.arr[this.arr.Length - 1]) + 1;

int low = 0, high = this.arr.Length - 1;

while (low < high)

{

int mid = low + (high - low) / 2;

if (arr[mid] == target)

{

return mid;

}

if (arr[mid] > target)

{

high = mid;

}

else

{

low = mid + 1;

}

}

return low;

}

Random rnd;

private int sum;

private int[] arr;

}

/**

* Your Solution object will be instantiated and called as such:

* Solution obj = new Solution(w);

* int param_1 = obj.PickIndex();

*/

Complexity: O(n) for construtor and O(logn) for pickIndex