Hex.*

LC528 - Random Pick With Weight

Posted on In
Leetcode: 136/2658 (57+65+14)

Description

You are given a 0-indexed array of positive integers w where w[i] describes the weight of the ith index.
You need to implement the function pickIndex(), which randomly picks an index in the range [0, w.length - 1] (inclusive) and returns it. The probability of picking an index i is w[i] / sum(w).
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%).

1
2
3
4
5
6
7
8
9
Example 1:
Input
["Solution","pickIndex"]
[[[1]],[]]
Output
[null,0]
Explanation
Solution solution = new Solution([1]);
solution.pickIndex(); // return 0. The only option is to return 0 since there is only one element in w.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
Example 2:
Input
["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
[[[1,3]],[],[],[],[],[]]
Output
[null,1,1,1,1,0]
Explanation
Solution solution = new Solution([1, 3]);
solution.pickIndex(); // return 1. It is returning the second element (index = 1) that has a probability of 3/4.
solution.pickIndex(); // return 1
solution.pickIndex(); // return 1
solution.pickIndex(); // return 1
solution.pickIndex(); // return 0. It is returning the first element (index = 0) that has a probability of 1/4.
Since this is a randomization problem, multiple answers are allowed.
All of the following outputs can be considered correct:
[null,1,1,1,1,0]
[null,1,1,1,1,1]
[null,1,1,1,0,0]
[null,1,1,1,0,1]
[null,1,0,1,0,0]
......
and so on.

Constraints:

1 <= w.length <= 104
1 <= w[i] <= 105
pickIndex will be called at most 104 times.

Solution

Selected by weight, here I am directly calling the random.randInt() function of Python. But this will only generate random numbers, not selected by weight. We can record the prefix sum, which is actually the sum of all numbers.

For example, the prefix sum of [1,3] is [1,4]. That is, if a random number of [1,4] is generated at this time, 1 corresponds to 1, and the random number 2,3,4 corresponds to the prefix and 4, which just satisfies the weight of 1:3.

And for w = [a,b,c,d] corresponding prefix and ws = [a,a+b,a+b+c,a+b+c+d] we can directly generate 0 to a+b+c +d-1 random number.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
# O(n) time | O(n) space
import random

class Solution:

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

    def pickIndex(self) -> int:
        val = random.randint(1, self.summ[-1])
        idx = 0
        for i, v in enumerate(self.summ):
            if v >= val:
                idx = i -1
                break
        return idx

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

Bisect Library

1
2
3
4
5
6
7
8
9
10
11
12
13
14
#import the module
import bisect

#given sorted list of numbers
nums = [1,3,5,7,10,25,49,55]

#given element to be inserted into the list
ele = 26

#get index where to insert the element
idx = bisect.bisect_left(nums, ele)

#print the index
print(f"Insert element {ele} at index {idx} in nums list to maintain sorted order.")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# O(n) time | O(n) space
class Solution:

    def __init__(self, w: List[int]):
        self.summ = [0]
        for i in w:
            self.summ.append(self.summ[-1]+i)
        print(self.summ)

    def pickIndex(self) -> int:
        num = random.randint(1, self.summ[-1])
        idx = bisect.bisect_left(self.summ, num)
        return idx-1

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