aboutsummaryrefslogtreecommitdiff
path: root/random_pick_weight.py
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())