Leetcode: 528. Random Pick with Weight

Patterns

Prompts

Can we pre-process the input in a way to make easy to solve the problem?: Compute the accumulate probability of picking any of \(0..i\) elements using Prefix Sum.
Can we formulate the problem as a classical problem?: Weighted Sample.

Solution

Compute the accumulate probability of picking any of \(0..i\) elements and use a Binary Search to find the index of one random element. Time complexity is \(O(n)\) for pre-processing and \(O(\log n)\) to answer each query. Space complexity is \(O(n)\) to store the accumulated probabilities.

class Solution:
    def __init__(self, w: List[int]):
        s = sum(w)
        self.a = list(accumulate([x / s for x in w]))

    def pickIndex(self) -> int:
        r = random.random()
        start, end, ans = 0, len(self.a) - 1, None
        while start <= end:
            mid = start + (end - start) // 2
            if r <= self.a[mid]:
                ans, end = mid, mid - 1
            else:
                start = mid + 1
        return ans