528. Random Pick with Weight

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

Python

Brute Force (Timelimit Exceed)

from random import randrange


class Solution:

    def __init__(self, w: List[int]):
        self.memory = []
        for index, weight in enumerate(w):
            for i in range(weight):
                self.memory.append(index)

    def pickIndex(self) -> int:
        pick = randrange(len(self.memory))
        return self.memory[pick]

Prefix Sum

• Reference: https://www.twblogs.net/a/5edebd5274efa30adcc735cc

from random import randrange


class Solution:

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

        self.prefix_sum = prefix_sum
        self.total = total

    def pickIndex(self) -> int:
        # target = 0 ~ total
        target = randrange(self.total)

        left, right = 0, len(self.prefix_sum)

        # Binary search the section of picked target
        while left < right:
            pivot = left + (right-left) // 2

            if target >= self.prefix_sum[pivot]:
                left = pivot + 1
            else:
                right = pivot

        return left