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 index0
is1 / (1 + 3) = 0.25
(i.e.,25%
), and the probability of picking index1
is3 / (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 most104
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