Randome Pick with Weight

class Solution:

    def __init__(self, w: List[int]):
        self.w = w
        self.prefixSum = []
        cumSum = 0
        for num in self.w:
            cumSum += num
            self.prefixSum.append(cumSum)


    def pickIndex(self) -> int:

        # imagine there is a stretched number line where a number
        # in w appears as many times as it is the value [1, 3] -> [1, 3, 3, 3]
        target = random.randint(0, sum(self.w) - 1)
        low, high = 0, len(self.w) - 1
        result = len(self.prefixSum) - 1

        while (low <= high):
            mid = (low + high) // 2

            if target < self.prefixSum[mid]:
                result = mid
                high = mid - 1
            else:
                low = mid + 1

        return result