[leetcode] 528. Random Pick with Weight @ python

原题

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.
Example 1:

Input:
[“Solution”,“pickIndex”]
[[[1]],[]]
Output: [null,0]
Example 2:

Input:
[“Solution”,“pickIndex”,“pickIndex”,“pickIndex”,“pickIndex”,“pickIndex”]
[[[1,3]],[],[],[],[],[]]
Output: [null,0,1,1,1,0]
Explanation of Input Syntax:

The input is two lists: the subroutines called and their arguments. Solution’s constructor has one argument, the array w. pickIndex has no arguments. Arguments are always wrapped with a list, even if there aren’t any.

解法

二分搜索法. 使用itertools.accumulate()构造列表储存w的累加的和, 使用bisect.bisect_left()返回将一个随机数插入数组时的index.
以[1,3]为例, 题意要求返回的index[0, 1]的概率应该是1:3. 我们构造的列表是[1,4], 然后产生随机数[1, 2, 3, 4], 当产生1时, 返回的index为0, 其他情况返回index1. 因此返回1的概率是返回0的概率的3倍.

代码

class Solution:

    def __init__(self, w: List[int]):
        self.data = list(itertools.accumulate(w))
        

    def pickIndex(self) -> int:
        return bisect.bisect_left(self.data, random.randint(1, self.data[-1]))
        


# Your Solution object will be instantiated and called as such:
# obj = Solution(w)
# param_1 = obj.pickIndex()