528. Random Pick with Weight

To do this, calculate the prefix sum of all the weights first. Then, we also calculate the sum total of all the weights. We then bisect to return the closest number to that and return it.

class Solution:
    def __init__(self, w: List[int]):
        # do a prefix sum and then return with that weight
        if not w:
            raise "Need at least one weight to pick"
        self.total = sum(w)
        self.weights = [w[0]]
        for weight in w[1:]:
            self.weights.append(self.weights[-1] + weight)
            
    def pickIndex(self) -> int:
        choice = random.randrange(0, self.total)
        picked = bisect.bisect_right(self.weights, choice)
        return picked