Solution
Understanding the question
The language used in the question is not the best and the example doesn’t further clarify the question. In simple terms, the question says that you have an array having some integers. Now, for each index i in the array, array[i] represents its weight. You have to pick up an index randomly while keeping in minding that higher weight should have higher probability. Suppose the array is [1,2] then index 1 should be twice as likely to be selected as index 2 as its weight is twice.
Simple approach [Time Limit Exceeded]
After understanding the question, the most basic approach which may come to your mind is creating another array where index i comes array[i] times and then simply picking a random number from that array. Suppose input array is [2,3,1] then we can translate it to an array where index 0 comes 2 times, index 1 comes 3 times and index 2 comes 1 time -> [0,0,1,1,1,2]. Picking a number randomly from this index will simply give us the answer.
This approach is correct is we have small weights and total size of array is small but the given question does not satisfy these constraints and creating such an array may give to TLE. Also, it wastes a huge amount of space.
Approach 2
Notice that the size of array we created in the above approach is essentially = sum(weights) and we can notice that index 0 occupies indices 0..weight[0]-1, index 1 occupies indices weight[0]..weight[1]-1, and index 2 occupies indices weight[1]..weight[2]-1.
This gives us a hint to preprocess the input array and store the prefix sum instead (arr[i] = arr[0]+arr[1]+…+arr[i]). So the above example becomes [2,5,6]. Now, if we choose any random number between 0 and 5 (inclusive), say r, and find the index of the number in the array which is just greater than r, we will have our answer. Suppose r=1, this gives us index 0 which is correct as in the array [0,0,1,1,1,2] above, index 1 (=r) has 0 stored. Suppose r=5, this gives us index 2 (6>5) which is correct as in the array [0,0,1,1,1,2] above, index 5 (=r) has 2 stored. To find this upper bound index, we can use binary search so our time complexity will be O(n + logn) = O(n).
The code below uses the STL function upper_bound which internally implements binary search:
class Solution {
public:
vector<int> v;
random_device rd;
Solution(vector<int>& w) {
v=w;
for (int i=1; i<v.size(); i++) {
v[i]+=v[i-1];
}
}
int pickIndex() {
int n = v[v.size()-1];
int random = rd()%n;
auto it = upper_bound(v.begin(), v.end(), random);
return it - v.begin();
}
};
Time complexity: O(n)