LeetCode 528 Random Pick with Weight

LeetCode

You are given a 0-indexed array of positive integers w where w[i] describes the weight of the i<sup>th</sup> 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 <= 10<sup>4</sup>
1 <= w[i] <= 10<sup>5</sup> pickIndex will be called at most 10<sup>4</sup> times.

题目大意：

根据数组每个元素的weight来决定其出现的概率： weight/sum of weight

解题思路：

模拟运算过程，先求和，然后根据上述公式分配概率： 如[1, 3], 小于0.25属于第一个元素，大于属于后一个元素，我们不用小数，还原回整数
所以数值小于1属于第一个元素，大于1小于4属于后一个，想到用presum，然后在presum搜索某个value，就想到二分法。

解题步骤：

N/A

注意事项：

1. random.randint前闭后闭

Python代码：

class Solution(TestCases):

    def __init__(self, w: List[int]):
        self.presum = []
        sum = 0
        for n in w:
            sum += n
            self.presum.append(sum)

    def pickIndex(self) -> int:
        rand_value = random.randint(0, self.presum[-1] - 1)
        return bisect.bisect(self.presum, rand_value)

算法分析：

时间复杂度为O(1ogn)，空间复杂度O(n)