Random pick with weight index

528.Random Pick with Weight

Binary Search

Prefix Sum

Math

Problem Statement:

Given an array w of positive integers, where w[i] describes the weight of index i, implement the function pickIndex() which randomly picks an index in proportion to its weight.

Example:

Input:

["Solution","pickIndex"] [[[1,3]],[]]

→

Output:

[null,1]

Java Solution:

// Extremeley easy problem, see solution for more info but heres a summary:
// Basically use preSum array and binary search of the keep searching varation
// 1. Get prefix sum array of weights
// 2. Generate random index from sum of all weights
// 3. Search for smallest index that is less than the random
class Solution {

    Random random;
    int[] wSums;  // Store prefix sums
    
    public Solution(int[] w) {
        this.random = new Random();
        for(int i=1; i// Build prefix sum array
        
        this.wSums = w;
    }
    
    public int pickIndex() {
        int len = wSums.length;
        int idx = random.nextInt(wSums[len-1]) + 1; // Random index from sum of all weights combined
        int left = 0, right = len - 1;
        int ans = 0;
        // search position
        while(left <= right){
            int mid = left + (right-left)/2;
            if(wSums[mid] == idx)
                return mid;
            else if (idx > wSums[mid])
                left = mid + 1;
            else {
               // If less, set ans and keep searching left
               ans = mid;
               right = mid - 1;
            }
        }
        return ans;
    }
}

Python Solution:

class Solution:
    def __init__(self, w: List[int]):
        # Build prefix sum array
        self.prefix = []
        prefix_sum = 0
        for weight in w:
            prefix_sum += weight
            self.prefix.append(prefix_sum)
        
    def pickIndex(self) -> int:
        target = random.randint(1, self.prefix[-1])
        left, right = 0, len(self.prefix) - 1
        
        # Binary search for target weight
        while left < right:
            mid = (left + right) // 2
            if self.prefix[mid] >= target:
                right = mid
            else:
                left = mid + 1
                
        return left

C++ Solution:

class Solution {
private:
    vector<int> prefix;
    
public:
    Solution(vector<int>& w) {
        // Build prefix sum array
        prefix = w;
        for(int i = 1; i < w.size(); i++)
            prefix[i] += prefix[i-1];
    }
    
    int pickIndex() {
        int target = rand() % prefix.back() + 1;
        return lower_bound(prefix.begin(), prefix.end(), target) - prefix.begin();
    }
};