题目
Given an array w of positive integers, where w[i] describes the weight of index i, write a function pickIndex which randomly picks an index in proportion to its weight.
Note:
- 1 <= w.length <= 10000
- 1 <= w[i] <= 10^5
- pickIndex will be called at most 10000 times.
解题思路
求石头重量的部分和,随机0到总和之间的一个随机数,二分搜索这个随机数的位置
代码
import random
import bisect
class Solution:
def __init__(self, w: List[int]):
self.sum = [w[0]]
for n in w[1:]:
self.sum.append(self.sum[-1] + n)
def pickIndex(self) -> int:
return bisect.bisect(self.sum, self.sum[-1] * random.random())
def pickIndex(self) -> int:
def search(n):
l, r = 0, len(self.sum)
while l < r:
m = (l + r) // 2
if self.sum[m] <= n:
l = m + 1
else:
r = m
return l
num = random.randint(0, self.sum[-1] - 1)
return search(num)
# Your Solution object will be instantiated and called as such:
# obj = Solution(w)
# param_1 = obj.pickIndex()