LC 0528 [M] Random Pick with Weight - ALawliet/algorithms GitHub Wiki

the probability of picking index i is w[i] / sum(w)
calculate the prefix sums and cum sum
randomly select a target by multiplying the cum sum with the input length
find the first prefix sum larger than target (binary search rightmost is faster than linear search)

For example, if w = [1, 3], the probability of picking index 0 is 1 / (1 + 3) = 0.25 (i.e., 25%),
and the probability of picking index 1 is 3 / (1 + 3) = 0.75 (i.e., 75%).

Explanation of why prefixSum works:

Think that if we had an array [1,2,3,4,3]. Normal random pickIndex would pick any index from 0 to 4 with equal probability. But we want that index=1 is picked by 2/13 probability, index=0 with 1/13 probability and so on. (13 is sum of weights). To ensure that one way to think of it if we make a larger array (of size 13) where the values are the indices such that index i is repeated w[i] times then if we do a normal rand on this array then index 0 to 12 will be picked randomly with equal probability. 13 index array -> [0, 1,1, 2,2,2, 3,3,3,3, 4,4,4]. So there is a 3/13 chance of picking 2 as 2 is repeated thrice in the new array.

Now instead of actually constructing this 13 index array, we just know the range of the index of the 13 index array where value = i. Eg:

for index=0, range is {0,0} index =1, range of indices of the new array is {1,2} index=2, range={3,5} index=3, range ={6,9} index = 4, range = {10,12} In other words,

index=0, range is <1 index=1, range is <3 index=2, range is <6 index = 3, range is < 10 index = 4, range is < 13 If you notice the above numbers 1,3,6,10,13 - they are cumulative sum. The reason this happens is because for every range: right = left + (w[i] - 1) and left is (prev right+1). So if we substitute 2nd equation into 1st. right = (prev right)+w[i]; i.e. keep adding prev sum to current weight.

Thus the prefixSum is able to implement this.

Alternate example to understand the solution:

So lets relate this problem with a problem of dividing a cake in a party such that the person with highest weight has better chance to pick a slice.(proportional to his weight)

Suppose we have people with weight 1, 3, 5, 7, 9 pounds and they went for a party to a bakery and purchased a cake. They decided that lets add our total weight and buy a cake proportional to it. So their totalWeight came as 1+3+5+7+9 = 25 They purchased a 25 pound cake :). They decided that lets cut cake in 1 pound slices(25 slices total) and whoever wants to eat can take a 1 pound slice at a time. The person who will pick a slice will be picked randomly.

To find out how to pick randomly and to figure out who will pick first, they first did a cumulative sum of their weights

1->1 (1, 1 slice)
3-> 1 + 3 = 4 (2-4, 3 slices)
5-> 4 + 5 = 9 (5-9, 5 slices)
7-> 7 + 9 = 16 (10-16, 7 slices)
9-> 9 + 16 = 25 (17-25, 9 slices)

=1,4,9,16,25

And then asked the server to pick a random number out of 25 for them. The random number represents a slice. So it can be 17 out of 25 or 6 out of 25.

So the slice 1 belongs to 1 pounds person. Which is only 1 slice. Slice between 2-4 belong to 3 pounds person. Which are 3 slices. . . Slice between 17- 25 belong to the 9 pounds person. Which are 9 slices.

If we pick a random slice out of 25, there is a higher probability that it will belong to a person with 9 slices (the person with 9 pounds) , the person who own 7 slices has second highest probability. The person whose weight is 1 pound and has only 1 slice has least probability to pick a slice.

And that's what we wanted. We want to let the person with highest weight get a greater chance to pick a slice every time even though they are picked at random.

The question can also be asked in reverse.

def __init__(self, w: List[int]):

self.prefixsums = list(accumulate(w))

self.prefixsums = [sum(w[:i+1]) for i in range(len(w))]

self.prefixsums = []
prefixsum = 0
for weight in w:
    prefixsum += weight
    self.prefixsums.append(prefixsum)

def pickIndex(self):

T = self.prefixsums[-1] * random.random() # totalsum * 0-1 decimal

# bisect == bisect_right, but in this case since the element exists, then bisect == bisect_right == bisect_left
l = bisect.bisect(self.prefixsums, T)

n = len(self.prefixsums)
l = 0
r = n
while l < r:
    m = (l + r) // 2
    if self.prefixsums[m] < T:
        l = m + 1
    else:
        r = m
return l

class Solution:

    def __init__(self, w: List[int]):
        self.prefixsums = list(accumulate(w))
            
    def pickIndex(self):
        totalsum = self.prefixsums[-1]
        x = totalsum * random.random() # totalsum * 0-1 decimal
        r = bisect_right(self.prefixsums, x) # first larger than target
        return r