Random Pick with Weight

Updated on 04 July, 2025
Random Pick with Weight header image

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 most 10^4 times.

Approach and Intuition

Understanding how to translate the given weights into a probabilistic selection of indices involves two main steps:

  1. 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 index 0, and numbers > 1 and <= 4 would map to index 1.

  2. Utilize Random Number for Index Selection: When pickIndex() is invoked, generate a random integer between 1 and the sum of w. 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:

  1. Calculate a prefix sum array prefixSums from w.

  2. Store the total sum totalSum = prefixSums[-1].

  3. On each pickIndex() call:

    • Generate a random number r in the range [1, totalSum].
    • Use binary search to find the index i such that prefixSums[i - 1] < r <= prefixSums[i].
    • Return the corresponding index i.

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++
cpp
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 array weights.
  • 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 multiplying randomValue with the last value in cumulativeWeights which is the total sum of weights.
    • Use lower_bound to find the first position in cumulativeWeights where the cumulative weight is greater than or equal to the threshold.
    • Return the position as the index, which corresponds to the weight segment that should be picked according to the random selection.

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
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 sum sumWeights.
  • 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:

  1. Initialize the array cumulativeWeights to store cumulative sum of the input weights.
  2. Iterate through the input weights. For each weight, add it to the cumulative sum from the previous iteration and update cumulativeWeights.
  3. Store the total weight in sumWeights.

Selecting an Index:

  1. Generate a random number scaled by the sum of all weights.
  2. Use binary search on the cumulativeWeights to find the appropriate index where the random number would fit.
  3. 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
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.

Comments

No comments yet.