LeetCode 528. Random Pick with Weight

December 10, 2020

We can do this in $O(n)$ time to initialize and $O(log(n))$ time to pick:

Build a list of left-to-right cumulative weights.
Randomly select a number, $n$ , $0 \leq n \leq sum(w)$ .
Do a binary search for the first cumulative weight $\geq n$ and return that index.

import bisect
import itertools
import random
from typing import List
import unittest

class Test(unittest.TestCase):

    def test_example_1(self):
        s = Solution([1])
        self.assertEqual(s.pickIndex(), 0)

    def test_example_2(self):
        s = Solution([1,3])
        self.assertIn(s.pickIndex(), {0,1})

class Solution:

    def __init__(self, w: List[int]):
        random.seed()
        self.sum = sum(w)
        self.cw = [x for x in itertools.accumulate(w)]

    def pickIndex(self) -> int:
        # This may or may not include <self.sum> in the
        # range (depending on floating point rounding)
        # so it might not be exactly correct.
        t = random.random() * self.sum
        return bisect.bisect_left(self.cw, t)

if __name__ == "__main__":
    unittest.main()