Here is the code:
import random
import bisect
class Solution:
def __init__(self, w):
self.prefix = []
total = 0
for weight in w:
total += weight
self.prefix.append(total)
self.total = total
def pickIndex(self):
r = random.uniform(0, self.total)
return bisect.bisect_right(self.prefix, r)
Time: O(n) for init, O(log n) for pickIndex. Space: O(n) for prefix array.