LeetCode 528. Random Pick with Weight - 二分查找（Binary Search）系列题8

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.

这题能否用二分查找法，一开始并不明朗，所以先来分析一下题目大意。 

题目给了一个长度为n的数组w代表有n个点（0~n-1），数组里的数值表示每个点的权重值(w[i]就是第i个点的权重值)，要求实现一个函数随机返回一个点，返回的概率是该点的权重值占总权重值的比例。

所有点的总权重值就是整个数组之和(totalWeight)，如果把[0, totalWeight)看成一个区间，那么每个点根据自己的权重值将占有[0, totalWeight)中一块子空间，如果把[0, totalWeight)区间按每个点的权重值划分成n个子区间, 那么所有子区间分布就是：[0, w[0]), [w[0], w[0] + w[1]),  [w[0] + w[1], w[0] + w[1] + w[2]), ... , 可以看出每个点所在区间的起点是该点前面所有点的权重和，终点是包括该点往前的所有点权重和也是下一个点所在区间的起点，因此所有点的子区间可以用一个数组来表示：[0, sumWeight[0], sumWeight[1], ..., sumWeight[n-1]] ( sumWeight[i]就是我们熟悉的数组前缀和)。如果从[0, totalWeight)随机取一个值，该值落在其中一个子区间的概率就是题目要求的概率。

(sumWeight[i] - sumWeight[i - 1])/totalWeight = w[i] / totalWeight

到此，本题就变成了从[0, totalWeight)里的随机取出一个值，看落在哪个子区间里从而求出子区间对应的是哪一个点。如前所述子区间是一个用权重数组的前缀和组成的数组来表示， 那要求随机值落在哪个子区间，就是前缀和数组里找出第一个比随机值大的点的前一个点，也就是求一个点在一个有序数组里的upper_bound。upper_bound函数就是用二分查找法实现的，在之前就已经遇到过了，代码实现相当简单。LeetCode 34. Find First and Last Position of Element in Sorted Array

 

class Solution:

    def __init__(self, w: List[int]):
        self.totalWeight = 0
        self.pos = []
        for i in w:
            self.pos.append(self.totalWeight)
            self.totalWeight += i

    def pickIndex(self) -> int:
        target = random.randint(0, self.totalWeight - 1)
        
        #upper bound
        l, r = 0, len(self.pos) - 1
        while l <= r:
            mid = l + (r - l) // 2
            if target  < self.pos[mid]:
                r = mid - 1
            else:
                l = mid + 1
        
        return l - 1