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).
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 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] <= 105pickIndex will be called at most 104 times.Problem Overview: You receive an array w where w[i] represents the weight of index i. The task is to implement pickIndex() so that the probability of picking index i is proportional to its weight. If w = [1,3], index 1 should be returned three times more often than index 0.
The core idea is to treat weights as ranges on a number line. If the total weight is 10, you randomly choose a number from [1,10] and determine which weight range contains that number.
Approach 1: Cumulative Sum with Linear Search (Initialization: O(n), Pick: O(n) time, O(n) space)
Build a prefix sum array where prefix[i] stores the total weight from index 0 to i. For example, weights [1,3,2] become prefix sums [1,4,6]. During pickIndex(), generate a random integer between 1 and the total weight. Iterate through the prefix array and return the first index where prefix[i] >= target. This works because each index occupies a proportional segment of the total range. The implementation is simple and demonstrates the key prefix sum idea, but every pick requires scanning the array.
Approach 2: Cumulative Sum with Binary Search (Initialization: O(n), Pick: O(log n) time, O(n) space)
The optimized approach keeps the same prefix sum array but replaces the linear scan with binary search. After generating a random number between 1 and the total weight, perform a lower-bound search on the prefix array to find the first index where the prefix value is greater than or equal to the target. Because the prefix array is sorted by construction, binary search finds the correct index in O(log n). This significantly improves performance when pickIndex() is called many times. The method combines array preprocessing, prefix sums, and randomized selection.
Recommended for interviews: The prefix sum with binary search solution is the expected answer. It shows that you recognize the weighted probability mapping and know how to convert it into a range search problem. Mentioning the linear scan approach first demonstrates understanding of the probability model, while the binary search optimization shows algorithmic maturity and awareness of repeated-query performance.
This approach involves creating a cumulative sum array based on the provided weights. The idea is to convert the weight array into a cumulative distribution, where each element represents the summed result of all previous weights including the current one. When we generate a random number, we search for its position in this cumulative array to determine which index to return.
This C solution implements a cumulative sum array ('prefixSum') to store the cumulative weights of the input array 'w'. The 'solutionPickIndex' function then generates a random number within the range of the total sum of weights and searches for the corresponding index in the cumulative sum array. The search is done linearly here.
Time Complexity: O(N) for initialization, O(N) for the pickIndex.
Space Complexity: O(N) due to the cumulative sum storage.
This optimized approach also uses a cumulative sum array, but instead of performing a linear search to find the appropriate index, we use a binary search. This greatly improves the efficiency when determining which index corresponds to a given cumulative value, especially beneficial for larger arrays.
This C implementation uses binary search to efficiently find the index corresponding to the random target. After setting up the cumulative sum array, a target number is generated, and binary search is used to pinpoint the index.
Time Complexity: O(N) for initialization, O(log N) for pickIndex.
Space Complexity: O(N) for the cumulative sum array.
| Approach | Complexity |
|---|---|
| Cumulative Sum with Linear Search | Time Complexity: O(N) for initialization, O(N) for the pickIndex. |
| Cumulative Sum with Binary Search | Time Complexity: O(N) for initialization, O(log N) for pickIndex. |
| Default Approach | — |
| Approach | Time | Space | When to Use |
|---|---|---|---|
| Cumulative Sum with Linear Search | Init: O(n), Pick: O(n) | O(n) | Simple implementation or when pickIndex is called only a few times |
| Cumulative Sum with Binary Search | Init: O(n), Pick: O(log n) | O(n) | Best choice when many random picks are required |
Random Pick with Weight | Bucket approach | Leetcode #528 | Binary search • Techdose • 41,902 views views
Watch 9 more video solutions →Practice Random Pick with Weight with our built-in code editor and test cases.
Practice on FleetCode