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).
This approach involves pre-calculating the cumulative sum of the weights. We can imagine the weights as contiguous blocks on a number line, where the size of each block corresponds to its weight. To pick an index, we generate a random number within the total range of this number line and then linearly scan through our pre-calculated cumulative sums to find which block (and thus which index) this random number falls into.
Algorithm
- In the constructor, create a
prefixSumsarray of the same size as the inputw. - Iterate through
w, calculating the cumulative sum at each indexi(w[0] + ... + w[i]) and storing it inprefixSums[i]. - Store the total sum, which is the last element of
prefixSums. - In the
pickIndexmethod, generate a random integertargetin the range[1, totalSum]. - Perform a linear search through the
prefixSumsarray. - Return the first index
ifor whichtarget <= prefixSums[i].
Walkthrough
Constructor (Solution(int[] w))
First, we need to preprocess the weights in the constructor to facilitate the picking process. We create a prefixSums array. prefixSums[i] will store the sum of all weights from index 0 up to i. The last element of this array will give us the total sum of all weights.
pickIndex() Method
When pickIndex() is called, we first generate a random integer, let's call it target, between 1 and the total sum (inclusive). Then, we iterate through the prefixSums array from the beginning. The first index i where target is less than or equal to prefixSums[i] is the index we are looking for. This works because the range of cumulative sum values corresponding to index i has a length of w[i], making the probability of picking i proportional to its weight.
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; // Linear search to find the index for (int i = 0; i < prefixSums.length; i++) { if (target <= prefixSums[i]) { return i; } } return -1; // Should not be reached }}Complexity
Time
- **Constructor:** O(N), where N is the length of `w`, to iterate through the array and build the prefix sums. - **`pickIndex()`:** O(N) in the worst case, as we might need to scan the entire `prefixSums` array to find the correct index.
Space
O(N), where N is the length of the input array `w`. This space is required to store the `prefixSums` array.
Trade-offs
Pros
The approach is relatively simple to understand and implement.
The preprocessing step in the constructor is efficient (linear time).
Cons
The
pickIndexoperation has a linear time complexity of O(N), which can be too slow and lead to a 'Time Limit Exceeded' error on large inputs or with many calls.
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(); *//** * @param {number[]} w */ var Solution = function (w) { const n = w.length; this.s = new Array(n + 1).fill(0); for (let i = 0; i < n; ++i) { this.s[i + 1] = this.s[i] + w[i]; }};/** * @return {number} */ Solution.prototype.pickIndex = function () { const n = this.s.length; const x = 1 + Math.floor(Math.random() * this.s[n - 1]); let left = 1, right = n - 1; while (left < right) { const mid = (left + right) >> 1; if (this.s[mid] >= x) { right = mid; } else { left = mid + 1; } } return left - 1;}; /** * Your Solution object will be instantiated and called as such: * var obj = new Solution(w) * var param_1 = obj.pickIndex() */class Solution { public: vector < int > s ; Solution ( vector < int >& w ) { int n = w . size (); s . resize ( n + 1 ); for ( int i = 0 ; i < n ; ++ i ) s [ i + 1 ] = s [ i ] + w [ i ]; } int pickIndex () { int n = s . size (); int x = 1 + rand () % s [ n - 1 ]; int left = 1 , right = n - 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(); */class Solution : def __init__ ( self , w : List [ int ]): self . s = [ 0 ] for c in w : self . s . append ( self . s [ - 1 ] + c ) def pickIndex ( self ) -> int : x = random . randint ( 1 , self . s [ - 1 ]) left , right = 1 , len ( self . s ) - 1 while left < right : mid = ( left + right ) >> 1 if self . s [ mid ] >= x : right = mid else : left = mid + 1 return left - 1 # Your Solution object will be instantiated and called as such: # obj = Solution(w) # 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.