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.The key idea in Random Pick with Weight is to ensure that each index is selected with probability proportional to its weight. A common strategy is to transform the weights into a prefix sum array. This array represents cumulative weight ranges, where each index occupies a segment proportional to its weight.
Once the prefix sums are prepared, generate a random number between 1 and the total weight. The problem then becomes finding which prefix range contains this number. Using binary search allows us to efficiently locate the corresponding index.
This approach works because larger weights create larger intervals in the prefix sum array, naturally increasing their selection probability. The preprocessing step builds the prefix sums once, while each random pick operation performs a binary search to locate the correct interval.
The preprocessing takes O(n) time, and each pick operation runs in O(log n) time with minimal extra space beyond the prefix array.
| Approach | Time Complexity | Space Complexity |
|---|---|---|
| Prefix Sum + Binary Search | Initialization: O(n), Pick: O(log n) | O(n) |
Techdose
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.
Time Complexity: O(N) for initialization, O(N) for the pickIndex.
Space Complexity: O(N) due to the cumulative sum storage.
1typedef struct {
2 int* prefixSum;
3 int n;
4} Solution;
5
6Solution* solutionCreate(int* w,
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.
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.
Time Complexity: O(N) for initialization, O(log N) for pickIndex.
Space Complexity: O(N) for the cumulative sum array.
1
public class Solution {
private int[] prefixSum;
private Random random;
public Solution(int[] w) {
prefixSum = new int[w.Length];
prefixSum[0] = w[0];
for (int i = 1; i < w.Length; i++) {
prefixSum[i] = prefixSum[i - 1] + w[i];
}
random = new Random();
}
public int PickIndex() {
int target = random.Next(prefixSum[prefixSum.Length - 1]);
int left = 0, right = prefixSum.Length - 1;
while (left < right) {
int mid = left + (right - left) / 2;
if (prefixSum[mid] <= target) left = mid + 1;
else right = mid;
}
return left;
}
}Watch expert explanations and walkthroughs
Practice problems asked by these companies to ace your technical interviews.
Explore More ProblemsJot down your thoughts, approach, and key learnings
Prefix sums convert individual weights into cumulative intervals. These intervals represent probability ranges, making it easy to select an index proportionally by generating a random number within the total weight range.
Yes, this problem or similar weighted random selection questions are common in technical interviews at large tech companies. It tests understanding of prefix sums, probability, and efficient searching techniques like binary search.
A prefix sum array is the most effective data structure for this problem. It allows quick mapping between random values and weighted ranges, enabling efficient selection using binary search.
The optimal approach uses a prefix sum array combined with binary search. Prefix sums convert weights into cumulative ranges, and a random number is mapped to one of these ranges using binary search to determine the correct index.
This C# implementation manually performs binary search to find the index with fewer comparisons, increasing efficiency in target determination.