3.23 Random Pick with Weight

3.23.1 Problem Metadata

3.23.2 Description

You are given a 0-indexed array of positive integers w where w[i] describes the weight of the i-th 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%).

3.23.3 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.

3.23.4 Constraints

  • 1 <= w.length <= 10^4
  • 1 <= w[i] <= 10^5
  • pickIndex will be called at most 10^4 times

3.23.5 Solution - Prefix Sum + Binary Search

3.23.5.1 Walkthrough

The key insight is to transform this weighted random selection problem into a range-finding problem using prefix sums.

Core Idea: - Assign each index a range proportional to its weight - Generate a random number in [1, totalWeight] - Find which index’s range contains this random number

Example: Given w = [1, 3, 2]

  1. Build prefix sum array: prefixSum = [1, 4, 6]
    • Index 0 owns range [1, 1] (weight = 1)
    • Index 1 owns range [2, 4] (weight = 3)
    • Index 2 owns range [5, 6] (weight = 2)
  2. Generate random number: Pick randVal uniformly in [1, totalWeight]
    • Example: randVal = 3
  3. Find corresponding index: Binary search in prefixSum to find where randVal belongs
    • randVal = 1 → falls in range [1, 1] → index 0
    • randVal = 2, 3, 4 → falls in range [2, 4] → index 1
    • randVal = 5, 6 → falls in range [5, 6] → index 2

Why This Works: - Each index i has w[i] possible random values that map to it - Total possible values = sum(w) - Probability of index i = w[i] / sum(w) ✓

Critical Detail: We use rand.nextInt(totalWeight) + 1 to generate values in [1, totalWeight] (not [0, totalWeight-1]), which aligns with our 1-indexed prefix sum ranges.

3.23.5.2 Analysis

Time Complexity: - Constructor: O(n) - Build prefix sum array with one pass - pickIndex(): O(log n) - Binary search in sorted array

Space Complexity: O(n) - Store prefix sum array

Why Binary Search over Linear Search?

The prefix sum array is monotonically increasing (always sorted), making it perfect for binary search:

  • Binary Search: O(log n) per call → With 10^4 calls: ~130,000 operations
  • Linear Search: O(n) per call → With 10^4 calls: ~100 million operations (worst case)

Performance difference: ~770x faster with binary search! As the array grows, linear search becomes unacceptable while binary search remains O(log n).

Handling Arrays.binarySearch() behavior: - Returns index ≥ 0 if exact match found - Returns -(insertion_point + 1) if not found (negative value) - We convert the negative result to get the insertion point: -(index + 1)

This insertion point is exactly the index we need, as it represents the first element ≥ randVal.

3.23.5.3 Implementation Steps

  1. Constructor - Precompute prefix sums:
    • Create prefixSum array of same length as w
    • Iterate through weights, accumulating cumulative sum
    • Store totalWeight as last element of prefix sum
  2. pickIndex() - Weighted random selection:
    • Generate random value: randVal = rand.nextInt(totalWeight) + 1 (range: [1, totalWeight])
    • Binary search prefixSum to find where randVal belongs
    • If exact match (index ≥ 0): return index directly
    • If not found (index < 0): convert to insertion point with -(index + 1)
  3. Edge cases handled:
    • Single weight: Always returns index 0
    • All equal weights: Uniform distribution
    • Large weights: Proper handling with int arithmetic

3.23.5.4 Code - Java

class Solution {
    private final int[] prefixSum;
    private final int totalWeight;
    private final Random rand = new Random();

    public Solution(int[] w) {
        int len = w.length;
        prefixSum  = new int[len];
        
        int curWeight = 0;
        for(int i = 0; i < len; i++) {
            int weight = w[i];
            curWeight += weight;
            prefixSum[i] = curWeight;
        }

        totalWeight = prefixSum[len -1];
    }
    
    public int pickIndex() {
        int randVal = rand.nextInt(totalWeight) + 1;

        int index = Arrays.binarySearch(prefixSum, randVal);

        // If exact match found, return it
        if (index >= 0) {
            return index;
        }
        
        // If not found, binarySearch returns -(insertion_point + 1)
        // We want the insertion point, which is where randVal would fit
        return -(index + 1);
    }
}