{x}
blog image

Random Pick with Weight

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 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 <= 104
  • 1 <= w[i] <= 105
  • pickIndex will be called at most 104 times.

Solution Explanation:

This problem requires generating random indices based on the weights provided in the input array w. The probability of selecting an index i is directly proportional to its weight w[i]. A naive approach would be to generate a random number and iterate through the weights, but this would be inefficient for large input arrays. The optimal solution utilizes prefix sums and binary search for efficient random index selection.

Algorithm:

  1. Prefix Sum Calculation: The __init__ (or constructor) method calculates the prefix sum of the weights. The prefix sum array s stores the cumulative sum of weights up to each index. For example, if w = [1, 3, 2], then s will be [0, 1, 4, 6]. s[i] represents the sum of weights from index 0 to i-1.

  2. Random Number Generation: The pickIndex method generates a random integer x between 1 and the total sum of weights (s[-1] or s[s.length - 1]). This ensures that the generated number falls within the range of the cumulative weights.

  3. Binary Search: A binary search is performed on the prefix sum array s to find the index i such that s[i] >= x. This index i is the randomly selected index. The binary search efficiently locates the index corresponding to the random number generated.

Time Complexity Analysis:

  • __init__ (Constructor): O(n), where n is the length of the input array w. This is due to the single pass required to compute the prefix sum.
  • pickIndex: O(log n). The binary search on the prefix sum array takes logarithmic time.

Space Complexity Analysis:

  • O(n) to store the prefix sum array s.

Code Explanation (Python Example):

import random
 
class Solution:
    def __init__(self, w: List[int]):
        self.s = [0]  # Initialize prefix sum array with 0
        for c in w:
            self.s.append(self.s[-1] + c)  # Calculate prefix sum
 
    def pickIndex(self) -> int:
        x = random.randint(1, self.s[-1])  # Generate random number within total weight range
        left, right = 1, len(self.s) - 1  # Set binary search boundaries
        while left < right:
            mid = (left + right) >> 1  # Efficient mid calculation using bitwise right shift
            if self.s[mid] >= x: #If the cumulative sum at mid is greater or equal to x, it means the index we are looking for is in the left half.
                right = mid  # Adjust right boundary
            else:
                left = mid + 1  # Adjust left boundary
        return left - 1  # Return the index

The other language examples follow the same logic, only differing in syntax and specific library functions (e.g., Random in Java, rand() in C++, rand.Intn in Go). The core algorithm remains consistent across all implementations.