random_pick_with_weight
1from bisect import bisect_left 2import random 3 4 5# @leet start 6class Solution: 7 """ 8 This question asks us to define a class, where, given a number of weights, 9 it randomly selects an index according to those weights. 10 11 One way to do this is to actually fill the range, where say given [1, 3] 12 you would create a weighted array of [0, 2, 2, 2] that corresponds to indexes 13 and then you would generate a random number within the range and that 14 would work. However, in the case that the numbers are very large, say 15 [1000000, 10000101, ...] we would need a very large array, and this would 16 become memory inefficient. A better way is to give up $O(1)$ random 17 generation in order to decrease the size of the array. 18 19 We calculate the prefix sum of all the weights and save that. Then, 20 we pick a random number from [0, 1] and then bisect left using that number 21 on the prefix sums. This allows us to find a random index in $O(log n)$ time, 22 but with an array that is the size of the weights $O(n)$, which is much faster 23 than the previous solution. 24 """ 25 26 def __init__(self, w: list[int]): 27 total = sum(w) 28 weights = [w[0] / total] 29 for weight in w[1:]: 30 weights.append(weights[-1] + weight / total) 31 self.weights = weights 32 33 def pickIndex(self) -> int: 34 random_num = random.random() 35 return bisect_left(self.weights, random_num) 36 37 38# Your Solution object will be instantiated and called as such: 39# obj = Solution(w) 40# param_1 = obj.pickIndex() 41# @leet end 42 43 44def test(): 45 assert 2 + 2 == 4
7class Solution: 8 """ 9 This question asks us to define a class, where, given a number of weights, 10 it randomly selects an index according to those weights. 11 12 One way to do this is to actually fill the range, where say given [1, 3] 13 you would create a weighted array of [0, 2, 2, 2] that corresponds to indexes 14 and then you would generate a random number within the range and that 15 would work. However, in the case that the numbers are very large, say 16 [1000000, 10000101, ...] we would need a very large array, and this would 17 become memory inefficient. A better way is to give up $O(1)$ random 18 generation in order to decrease the size of the array. 19 20 We calculate the prefix sum of all the weights and save that. Then, 21 we pick a random number from [0, 1] and then bisect left using that number 22 on the prefix sums. This allows us to find a random index in $O(log n)$ time, 23 but with an array that is the size of the weights $O(n)$, which is much faster 24 than the previous solution. 25 """ 26 27 def __init__(self, w: list[int]): 28 total = sum(w) 29 weights = [w[0] / total] 30 for weight in w[1:]: 31 weights.append(weights[-1] + weight / total) 32 self.weights = weights 33 34 def pickIndex(self) -> int: 35 random_num = random.random() 36 return bisect_left(self.weights, random_num)
This question asks us to define a class, where, given a number of weights, it randomly selects an index according to those weights.
One way to do this is to actually fill the range, where say given [1, 3] you would create a weighted array of [0, 2, 2, 2] that corresponds to indexes and then you would generate a random number within the range and that would work. However, in the case that the numbers are very large, say [1000000, 10000101, ...] we would need a very large array, and this would become memory inefficient. A better way is to give up $O(1)$ random generation in order to decrease the size of the array.
We calculate the prefix sum of all the weights and save that. Then, we pick a random number from [0, 1] and then bisect left using that number on the prefix sums. This allows us to find a random index in $O(log n)$ time, but with an array that is the size of the weights $O(n)$, which is much faster than the previous solution.