528. Random Pick with Weight

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

TLE Solution

import random
class Solution:

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

    def pickIndex(self) -> int:
        randIndex = random.randint(0,len(self.nums)-1)
        return self.nums[randIndex]
        

Instead of creating w[i] items for each index i, we use presum weight so we don't expand the array too large

w = [1,2,3,4]
presumW= [1,3,6,10]
if random index in
0 - 1 ====> 0
1 - 3 ====> 1
3 - 6 ====> 2
6 - 10 ===> 3

then it is same as find binary search left bound

for example: 1,3,4,6,10

4 falls in range(3,6) which is index 2

import random
class Solution:

    def __init__(self, w: List[int]):
        self.presumW = [0] * len(w)
        self.presumW[0]= w[0]
        for i in range(1,len(w)):
            self.presumW[i] = self.presumW[i-1] + w[i]
        
            
        

    def pickIndex(self) -> int:
        randIndex = random.randint(1,self.presumW[-1])
        left = 0
        right = len(self.presumW)-1
        while left < right:
            mid = (right-left)//2 + left
            if self.presumW[mid] == randIndex:
                return mid
            elif self.presumW[mid] < randIndex:
                left = mid+1
            else:
                right = mid
        return left