528. Random Pick with Weight
按权重随机选择
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 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 <= 10
4
1 <= w[i] <= 10
5
pickIndex
will be called at most10
4
times.
分析:
通过前缀和数组 + 二分查找实现按权重随机选择,确保每个下标被选中的概率与其权重成正比。前缀和数组将权重映射为连续区间,每个下标 i
对应区间 [prefix_sum[i-1] + 1, prefix_sum[i]]
。
import random
from itertools import accumulate
class Solution:
def __init__(self, w: List[int]):
self.presum = list(accumulate(w))
self.w = w
def pickIndex(self) -> int:
i = random.randint(1, self.presum[-1])
return bisect.bisect_left(self.presum, i)
# Your Solution object will be instantiated and called as such:
# obj = Solution(w)
# param_1 = obj.pickIndex()
Last updated
Was this helpful?