Random Pick with Weight

Intuition

Solution

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

    def pickIndex(self) -> int:
        value = random.randint(0, self.w[-1] - 1)
        if value < self.w[0]:
            return 0
        lo = 0
        hi = len(self.w) - 1
        while hi - lo > 1:
            mid = (lo + hi) // 2
            if self.w[mid] <= value:
                lo = mid
            else:
                hi = mid
        return hi


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