Random Pick with Weight

Published: Dec 28, 2022

Problem 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%).

Constraints:

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

https://leetcode.com/problems/random-pick-with-weight/

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]

How to Solve

This problem requires some ideas. The prefix sum and binary search are a good approach. The prefix sum reflects the difference of each weight. Do binary search on the prefix sum to get the index of a randomly generated weight. It depends on the languages what values are generated by random function. For example, C++ doesn’t need to regularize the given weights, but others need to change values between 0 to 1.

Solution

class RandomPickWithWeight {
private:
    vector<int> weights;
public:
    Solution(vector<int>& w) {
        srand(time(NULL));
        weights.push_back(w[0]);
        for (int i = 1; i < w.size(); ++i) {
            weights.push_back(w[i] + weights.back());
        }
    }
    
    int pickIndex() {
        int weight = rand() % weights.back();
        int left = 0, right = weights.size() - 1;
        int mid;
        while (left < right) {
            mid = (left + right) / 2;
            if (weight >= weights[mid]) {
                left = mid + 1;
            } else {
                right = mid;
            }
        }
        return left;
    }
};

class RandomPickWithWeight:

    def __init__(self, w: List[int]):
        total = sum(w)
        self.weights = [w[0] / total]
        for v in w[1:len(w) - 1]:
            self.weights.append(self.weights[-1] + v / total)
        self.weights.append(1)

    def pickIndex(self) -> int:
        weight = random.random()
        left, right = 0, len(self.weights) - 1
        while left < right:
            mid = (left + right) // 2
            if weight >= self.weights[mid]:
                left = mid + 1
            else:
                right = mid
        return left

class RandomPickWithWeight

=begin
    :type w: Integer[]
=end
    def initialize(w)
        total = w.sum
        @weights = [w[0].to_f / total]
        (1...w.size - 1).each do |i|
            @weights << w[i].to_f / total + @weights[-1]
        end
        @weights << 1
    end


=begin
    :rtype: Integer
=end
    def pick_index()
        weight = rand
        left, right = 0, @weights.size - 1
        while left < right
            mid = (left + right) / 2
            if weight >= @weights[mid]
                left = mid + 1
            else
                right = mid
            end
        end
        left
    end
end

Complexities

• Time: constructor: O(n), pickIndex: O(log(n))
• Space: O(n)