528. Random Pick with Weight

https://leetcode.com/problems/random-pick-with-weight/

solution

# 先构造一个累加的概率数组(prefix sum)，然后用二分查找
class Solution:
    def __init__(self, w):
        self.w = list(itertools.accumulate(w))

    def pickIndex(self):
        rand = random.randint(1, self.w[-1])  # 注意是1 ~ prefix sum的范围
        return bisect.bisect_left(self.w, rand)

时间复杂度：O(n) / O(log(n)) 空间复杂度：O(n)

class Solution:
    def __init__(self, w: List[int]):
        self.accumulate_list = [w[0]]
        for i in range(1, len(w)):
            self.accumulate_list.append(self.accumulate_list[-1] + w[i])

    def pickIndex(self) -> int:
        # 注意rand从1开始选，但下面的l作为index仍然从0开始
        rand = random.randint(1, self.accumulate_list[-1])

        l = 0
        r = len(self.accumulate_list)

        while l < r:
            mid = l + (r - l) // 2
            if self.accumulate_list[mid] < rand:
                l = mid + 1
            else:  # 第一个比rand大的index，等于时也shrink右边界
                r = mid
        return l