Random Pick with Weight
MedPrompt
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 index0is1 / (1 + 3) = 0.25(i.e.,25%), and the probability of picking index1is3 / (1 + 3) = 0.75(i.e.,75%).
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 <= 1041 <= w[i] <= 105pickIndexwill be called at most104times.
Approaches
2 approaches with complexity analysis and trade-offs.
This approach significantly optimizes the picking process. While the setup in the constructor remains the same—calculating prefix sums—it leverages a key property of the prefixSums array: it is sorted in non-decreasing order. This allows us to use a much faster binary search algorithm, instead of a linear scan, to find the correct index for a given random value. This reduces the time complexity of each pickIndex call from linear to logarithmic.
Algorithm
- The constructor logic is the same: create a
prefixSumsarray by calculating the cumulative sums of weights. - In
pickIndex, generate a random integertargetfrom 1 tototalSum. - Instead of a linear scan, perform a binary search on the
prefixSumsarray to find the target index. - Initialize
low = 0,high = prefixSums.length - 1, and aresultvariable. - While
low <= high, calculate the middle indexmid. - If
prefixSums[mid] >= target,midis a potential answer. Store it inresultand search for a potentially smaller index in the left half by settinghigh = mid - 1. - Otherwise, if
prefixSums[mid] < target, the answer must be in the right half, so setlow = mid + 1. - After the loop,
resultwill hold the smallest indexisuch thatprefixSums[i] >= target.
Walkthrough
Constructor (Solution(int[] w))
The constructor implementation is identical to the linear search approach. We compute a prefixSums array where prefixSums[i] holds the sum of weights from w[0] to w[i], and we store the totalSum.
pickIndex() Method
The key improvement is in this method. After generating a random target between 1 and totalSum, we use binary search to find the index. We are looking for the smallest index i such that prefixSums[i] >= target. This is a classic 'lower bound' search problem.
The binary search works by repeatedly dividing the search interval in half. If the value at the middle of the interval is less than the target, we know the target must be in the right half. If the value is greater than or equal to the target, the middle element could be our answer, but we continue searching in the left half to see if a smaller index also satisfies the condition. This process efficiently hones in on the correct index.
class Solution { private int[] prefixSums; private int totalSum; private java.util.Random rand = new java.util.Random(); public Solution(int[] w) { this.prefixSums = new int[w.length]; int currentSum = 0; for (int i = 0; i < w.length; ++i) { currentSum += w[i]; this.prefixSums[i] = currentSum; } this.totalSum = currentSum; } public int pickIndex() { // Generate a random number between 1 and totalSum (inclusive) int target = rand.nextInt(totalSum) + 1; // Binary search to find the index int low = 0; int high = prefixSums.length - 1; while (low < high) { int mid = low + (high - low) / 2; if (target > prefixSums[mid]) { low = mid + 1; } else { high = mid; } } return low; }}Complexity
Time
- **Constructor:** O(N), where N is the length of `w`, for the one-time setup of the prefix sum array. - **`pickIndex()`:** O(log N) due to the efficient binary search on the `prefixSums` array.
Space
O(N), where N is the length of the input array `w`. This space is used to store the `prefixSums` array.
Trade-offs
Pros
Extremely efficient
pickIndexoperation with O(log N) time complexity, making it ideal for frequent calls.It's the optimal approach for the given problem constraints.
Cons
Requires O(N) extra space for the prefix sums, which might be a concern for memory-constrained environments with extremely large inputs (though it's acceptable for the given constraints).
Solutions
Solution
class Solution { private int [] s ; private Random random = new Random (); public Solution ( int [] w ) { int n = w . length ; s = new int [ n + 1 ]; for ( int i = 0 ; i < n ; ++ i ) { s [ i + 1 ] = s [ i ] + w [ i ]; } } public int pickIndex () { int x = 1 + random . nextInt ( s [ s . length - 1 ]); int left = 1 , right = s . length - 1 ; while ( left < right ) { int mid = ( left + right ) >> 1 ; if ( s [ mid ] >= x ) { right = mid ; } else { left = mid + 1 ; } } return left - 1 ; } } /** * Your Solution object will be instantiated and called as such: * Solution obj = new Solution(w); * int param_1 = obj.pickIndex(); */Video walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.