528. Random Pick with Weight - cocoder39/coco39_LC GitHub Wiki

528. Random Pick with Weight

attention to use random.randint(1, self.sum) instead of random.randint(0, self.sum). every number is positive. say w = [1, 3], then total weight is 4 and the possibility of 1 is 1/4. if we do random.randint(0, self.sum), then total number of positive is 5 (from 0 to 5). if we say both 0 and 1 fall into index 0 then the p is 2/5 instead of 1/4

class Solution:

    def __init__(self, w: List[int]):
        self.prefix_sum = []
        self.sum = 0
        for i, weight in enumerate(w):
            self.sum += weight
            self.prefix_sum.append(self.sum)

    def pickIndex(self) -> int:
        def helper(num):
            low, high = 0, len(self.prefix_sum) - 1
            while low + 1 < high:
                mid = low + (high - low) // 2
                if self.prefix_sum[mid] == num:
                    return mid
                elif self.prefix_sum[mid] > num:
                    high = mid
                else:
                    low = mid
            
            if self.prefix_sum[low] >= num:
                return low
            return high
                    
        num = random.randint(1, self.sum)
        return helper(num)