blob: 3eda7fdb7b9641032ed6ba0478e67695195e971c (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
|
# 528. Random Pick with Weight
from itertools import accumulate
from bisect import bisect_left
from random import randint
"""
you are given a 0-indexed array of positive integers 'w' where 'w[i]'
describes the weight of the i'th index. you need to implement the function
'pickIndex()' which randomly picks an index in the range '[0, w.len - 1]' and
returns it. the probability of picking an index 'i' is 'w[i] / sum(w)'.
"""
class Solution(object):
def __init__(self, w):
"""
:type w: List[int]
"""
self.w = list(accumulate(w))
def pickIndex(self):
"""
:rtype: int
"""
return bisect_left(self.w, randint(1, self.w[-1]))
if __name__ == "__main__":
obj1 = Solution([1])
obj2 = Solution([1, 3])
print(obj1.pickIndex())
print(obj1.pickIndex())
print(obj1.pickIndex())
print(obj1.pickIndex())
print(obj1.pickIndex())
print(obj1.pickIndex())
|