We can binary-search in the cumulative weight list to pick an index with frequency proportional to its weight.

| Operation | Time | Space |
__init__ | ||
pickIndex |
class Solution:
def __init__(self, w: List[int]):
self.cum_weight = [0] * len(w)
self.cum_weight[0] = w[0]
for i in range(1, len(w)):
self.cum_weight[i] = self.cum_weight[i-1] + w[i]
def __find_insert_index(self, draw):
lo, hi = 0, len(self.cum_weight)-1
while lo <= hi:
mid = (lo+hi) // 2
if self.cum_weight[mid] < draw:
lo = mid+1
else:
hi = mid-1
return lo
def pickIndex(self) -> int:
draw = randint(1, self.cum_weight[-1])
return self.__find_insert_index(draw)
# Your Solution object will be instantiated and called as such:
# obj = Solution(w)
# param_1 = obj.pickIndex()
Leave a comment