
Problem Statement
In this task, you're working with an array of positive integers w
, where each element w[i]
represents the weight associated with the i-th
index. Your objective is to implement a function named pickIndex()
. This function selects an index between 0
and w.length - 1
, inclusive. However, it doesn't just pick any index uniformly; instead, it selects an index based on a probabilistic model where the likelihood of choosing an index i
is proportional to w[i]
. Specifically, the probability of picking index i
is given by the fraction w[i] / sum(w)
, where sum(w)
is the sum of all weights in the array.
For instance, consider the array w = [1, 3]
. Here, the probability for index 0
would be 1/4 (or 25%)
and for index 1
would be 3/4 (or 75%)
. The function pickIndex()
should reflect this probability distribution when called, meaning index 1
should roughly be picked three times more frequently than index 0
.
Examples
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 <= 10^4
1 <= w[i] <= 10^5
pickIndex
will be called at most10^4
times.
Approach and Intuition
Understanding how to translate the given weights into a probabilistic selection of indices involves two main steps:
Calculate Cumulative Sums: Convert the array of weights into an array of cumulative sums. This transformation allows for an efficient way of mapping a uniformly generated random number to the correct index based on weight. For example, with
w = [1, 3]
, the cumulative array would be[1, 4]
. Here, numbers<= 1
would map to index0
, and numbers> 1 and <= 4
would map to index1
.Utilize Random Number for Index Selection: When
pickIndex()
is invoked, generate a random integer between1
and the sum ofw
. Then perform a binary search on the cumulative array to find the smallest index such that the cumulative value is greater than or equal to the random number. This index corresponds to the weighted random pick.
Detailed Steps:
Calculate a prefix sum array
prefixSums
fromw
.Store the total sum
totalSum = prefixSums[-1]
.On each
pickIndex()
call:- Generate a random number
r
in the range[1, totalSum]
. - Use binary search to find the index
i
such thatprefixSums[i - 1] < r <= prefixSums[i]
. - Return the corresponding index
i
.
- Generate a random number
This approach ensures that each index is picked with a probability proportional to its weight, leveraging the prefix sum structure and binary search for optimal performance.
Solutions
- C++
class Solution {
vector<int> cumulativeWeights;
public:
Solution(vector<int> &weights) {
for (auto weight : weights)
cumulativeWeights.push_back(weight + (cumulativeWeights.empty() ?
0 : cumulativeWeights.back()));
}
int pickIndex() {
float randomValue = (float) rand() / RAND_MAX;
float threshold = randomValue * cumulativeWeights.back();
return lower_bound(begin(cumulativeWeights), end(cumulativeWeights), threshold) - begin(cumulativeWeights);
}
};
The provided C++ solution efficiently handles the task of selecting an index based on given weights. The algorithm works by constructing a cumulative distribution of the weights which allows for an efficient selection process using binary search. The implementation follows these detailed steps:
- Initialize a class Solution with a vector member
cumulativeWeights
that stores the cumulative sum of the input arrayweights
. - In the constructor of the Solution class, compute the cumulative weights. For each weight in the 'weights' vector, add it to the last value in the
cumulativeWeights
vector if this is not the first weight; otherwise, start with the weight itself. - The
pickIndex
function picks an index randomly according to the weight distribution:- Generate a random float value
randomValue
between 0 and 1. - Calculate a
threshold
by multiplyingrandomValue
with the last value incumulativeWeights
which is the total sum of weights. - Use
lower_bound
to find the first position incumulativeWeights
where the cumulative weight is greater than or equal to thethreshold
. - Return the position as the index, which corresponds to the weight segment that should be picked according to the random selection.
- Generate a random float value
This approach guarantees that the selection of indices is proportional to their weights, while maintaining efficient query time through the use of cumulative sums and binary search for index picking.
- Java
class RandomPicker {
private int[] cumulativeWeights;
private int sumWeights;
public RandomPicker(int[] weights) {
cumulativeWeights = new int[weights.length];
int cumulativeSum = 0;
for (int i = 0; i < weights.length; i++) {
cumulativeSum += weights[i];
cumulativeWeights[i] = cumulativeSum;
}
sumWeights = cumulativeSum;
}
public int selectIndex() {
double randomTarget = sumWeights * Math.random();
// Binary search to find the right index
int left = 0, right = cumulativeWeights.length;
while (left < right) {
int middle = left + (right - left) / 2;
if (randomTarget > cumulativeWeights[middle])
left = middle + 1;
else
right = middle;
}
return left;
}
}
The problem presents a scenario where you need to randomly pick an index from a list, given varying probabilities denoted by the weight of each index. The solution is implemented using Java and involves constructing a RandomPicker
class.
Overview:
- The class maintains a cumulative weight sum using an array
cumulativeWeights
and a total sumsumWeights
. - Constructor
RandomPicker
initializes the cumulative weights array based on the input weights. It calculates a running total (cumulative sum) of the weights to set up for efficient index picking.
Steps Involved:
- Initialize the array
cumulativeWeights
to store cumulative sum of the input weights. - Iterate through the input weights. For each weight, add it to the cumulative sum from the previous iteration and update
cumulativeWeights
. - Store the total weight in
sumWeights
.
Selecting an Index:
- Generate a random number scaled by the sum of all weights.
- Use binary search on the
cumulativeWeights
to find the appropriate index where the random number would fit. - The
selectIndex
method effectively maps the random number to an index based on its weighted probability.
Key Points:
- The use of cumulative sums allows the
selectIndex
method to leverage binary search, optimizing the selection operation to O(log n) complexity. - The random index reflects the probability distribution given by the initial weights, allowing for a fair selection based on the provided weights.
This solution ensures that the performance is efficient even with a large input size, and it adheres to the weighted probability distribution accurately.
- Python
class Solution:
def __init__(self, weights: List[int]):
self.accumulated_weights = []
accumulated = 0
for weight in weights:
accumulated += weight
self.accumulated_weights.append(accumulated)
self.total_weight = accumulated
def pickIndex(self) -> int:
target = self.total_weight * random.random()
low, high = 0, len(self.accumulated_weights)
while low < high:
mid = (low + high) // 2
if target > self.accumulated_weights[mid]:
low = mid + 1
else:
high = mid
return low
The provided Python solution implements a class for selecting a random index based on weights given to constructor. The Solution
class has two methods:
__init__(self, weights: List[int])
: Initializes the object using the weights passed to it. It creates a list,self.accumulated_weights
, that stores the cumulative sum of the weights. This cumulative sum aids in efficiently picking an index based on weight.pickIndex(self) -> int
: Returns a random index, weighted according to the initial list passed to the constructor. This method uses binary search to select an index based on its probability weight. Here is how the method works:- Generate a random number scaled to the sum of weights.
- Use binary search on the accumulated weights to find the corresponding index for this random number.
This method ensures that an index with higher weight is more likely to be picked. Overall, the code is efficient with initialization done in (O(n)) time complexity where (n) is the length of the input list, and the pickIndex
function operates in (O(\log n)) time complexity due to the binary search implementation.
No comments yet.