[LeetCode] 528. Random Pick with Weight

You are given an array of positive integers w where w[i] describes the weight of ith index (0-indexed).

We need to call the function pickIndex() which randomly returns an integer in the range [0, w.length - 1]. pickIndex() should return the integer proportional to its weight in the w array. For example, for w = [1, 3], the probability of picking the index 0 is 1 / (1 + 3) = 0.25 (i.e 25%) while the probability of picking the index 1 is 3 / (1 + 3) = 0.75 (i.e 75%).

More formally, the probability of picking index i is w[i] / sum(w). 

Example 1:

Input
["Solution","pickIndex"]
[[[1]],[]]
Output
[null,0]

Explanation
Solution solution = new Solution([1]);
solution.pickIndex(); // return 0. Since there is only one single element on the array the only option is to return the first element.

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's returning the second element (index = 1) that has probability of 3/4.
solution.pickIndex(); // return 1
solution.pickIndex(); // return 1
solution.pickIndex(); // return 1
solution.pickIndex(); // return 0. It's returning the first element (index = 0) that has probability of 1/4.

Since this is a randomization problem, multiple answers are allowed so 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 <= 10000
• 1 <= w[i] <= 10^5
• pickIndex will be called at most 10000 times.

按权重随机选择。

给你一个 下标从 0 开始 的正整数数组 w ，其中 w[i] 代表第 i 个下标的权重。

请你实现一个函数 pickIndex ，它可以 随机地 从范围 [0, w.length - 1] 内（含 0 和 w.length - 1）选出并返回一个下标。选取下标 i 的 概率 为 w[i] / sum(w) 。

例如，对于 w = [1, 3]，挑选下标 0 的概率为 1 / (1 + 3) = 0.25 （即，25%），而选取下标 1 的概率为 3 / (1 + 3) = 0.75（即，75%）。

来源：力扣（LeetCode）
链接：https://leetcode.cn/problems/random-pick-with-weight
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

注意这个题最后返回的是权重w里面的index。思路是前缀和 + 二分法。我直接用例子讲，比如权重w = [2, 5, 3, 4], 那么权重和wsum = [2, 7, 10, 14]。权重和其实就是前缀和。

对于w中的每一个元素，被抽中的几率分别是 [2/14, 5/14, 3/14, 4/14]。

因为整个input的权重和是14，所以取random数字可以在[1, 14]这个范围内取，然后可以把这14个数字按照权重分配给w的下标，结果如下

idx in [1,2] return 0
idx in [3,7] return 1
idx in [8,10] return 2
idx in [11,14] return 3

当得到这个随机数之后，就去权重和里面找，比如如果随机数是4，在wsum = [2, 7, 10, 14]范围内做二分。因为 wsum[i] 一定是某个区间的上界，所以在做二分的时候，如果 sum[mid] < idx 那么left = mid + 1，说明idx肯定不在以 sum[mid] 为上界的这一段里面。按照这个思路跑，这个例子最后返回的idx应该是1。

时间O(n)

空间O(n)

Java实现

class Solution {
    Random random;
    int[] sum;

    public Solution(int[] w) {
        this.random = new Random();
        for (int i = 1; i < w.length; i++) {
            w[i] += w[i - 1];
        }
        this.sum = w;
    }

    public int pickIndex() {
        int len = sum.length;
        // 因为random是从[0, len)取值的所以要 + 1
        int idx = random.nextInt(sum[len - 1]) + 1;
        int left = 0;
        int right = len - 1;
        // search position 
        while (left < right) {
            int mid = left + (right - left) / 2;
            if (sum[mid] == idx) {
                return mid;
            } else if (sum[mid] < idx) {
                left = mid + 1;
            } else {
                right = mid;
            }
        }
        return left;
    }
}

 

LeetCode 题目总结