LeetCode 528. Random Pick with Weight

23 Aug 2018

Description:

Given an array w of positive integers, where w[i] describes the weight of index i, write a function pickIndex which randomly picks an index in proportion to its weight.

Note:

1 <= w.length <= 10000
1 <= w[i] <= 10^5

pickIndex will be called at most 10000 times.

Example:

input:
[“Solution”,”pickIndex”]
[[[1]],[]]
output:

[null,0,1,1,1,0]

input:

[“Solution”,”pickIndex”,”pickIndex”,”pickIndex”,”pickIndex”,”pickIndex”]
[[[1,3]],[],[],[],[],[]]

output:

[null,0,1,1,1,0]

Explanation of Input Syntax:

The input is two lists: the subroutines called and their arguments. Solution’s constructor has one argument, the array w. pickIndex has no arguments. Arguments are always wrapped with a list, even if there aren’t any.

Solution:

class Solution:

    pickLst = None
    
    def __init__(self, w):
        """
        :type w: List[int]
        """
        self.pickLst = [0.0]
        
        sum = 0
        for ele in w:
            sum += ele
        
        for ele in w:
            self.pickLst.append(self.pickLst[-1] + ele / sum)

    def pickIndex(self):
        """
        :rtype: int
        """
        rnd = random.random()
        
        left = 0
        right = len(self.pickLst) - 1
        result = -1
        
        while(right > left):
            mid = math.ceil((left + right) / 2)
            if self.pickLst[mid - 1] <= rnd < self.pickLst[mid]:
                result = mid - 1
                break
            if rnd < self.pickLst[mid - 1]:
                right = mid
            elif rnd > self.pickLst[mid]:
                left = mid
                
        if result == -1:
            raise RuntimeError('cannot find')
            
        return result


# Your Solution object will be instantiated and called as such:
# obj = Solution(w)
# param_1 = obj.pickIndex()

This is simulation of multinomial distribution. We can use bins and random floating point numbers to achieve simulation.

Alan Xiang's Blog