Random Pick with Weight
Problem Description
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
Example 2:
Input
["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
[[[1,3]],[],[],[],[],[]]
Output
[null,1,1,1,1,0]
Explanation
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 <= 104
- 1 <= w[i] <= 105
pickIndexwill be called at most104times.
Solve it on LeetCode
Approaches
1. Prefix Sum + Binary Search Approach
Intuition:
The main goal is to pick an index i with a probability proportional to w[i]. If we imagine each weight as a segment on a line of length sum(weights), then we can pick a random point on this line and determine which segment it falls into using binary search. Thus, the problem boils down to constructing a prefix sum array and using binary search to determine our pick.
- Compute prefix sums of the weights. This helps in cumulatively understanding the "length" of each segment.
- Use random function to pick a target in the range of 0 to the total sum of weights.
- Use binary search to find the index where the target falls in the prefix sum array.
Code:
- Time Complexity: Preprocessing: O(N), where N is the length of the weights array to compute prefix sums. Pick Index: O(log N) due to binary search.
- Space Complexity: O(N) for storing the prefix sum array.
2. Binary Search Tree Approach (Optimal)
Intuition:
This approach uses a dynamic data structure like a balanced binary search tree (e.g., TreeMap in Java) to store cumulative weights, allowing for efficient cumulative lookup and random access. This is especially efficient if weights frequently update or if the number of weights (N) is very large.
- Store cumulative weights in a
TreeMap, where each key is a cumulative weight, and the value is the corresponding index. - For each
pickIndexcall, randomly determine a target and find the smallest key in the map greater or equal to the target.
Code:
- Time Complexity: Preprocessing: O(N log N) due to insertions into the tree map. Pick Index: O(log N) since searching through the
TreeMapis logarithmic on average. - Space Complexity: O(N) due to storing the cumulative sums in the
TreeMap.