random pick with weight

🏠
 1 from random import randrange
 2 from itertools import accumulate
 3 
 4 class Solution:
 5 
 6     def __init__(self, w):
 7         self.w = w
 8         self.s = sum(self.w)
 9         acc = accumulate(self.w)
10         self.enum = []
11         enum = [*enumerate(acc)]
12         for i, a in enum:
13             l = enum[i - 1][1] if i > 0 else 0
14             u = a if i > 0 else self.w[0]
15             self.enum.append((l, u, i))
16 
17     def pickIndex(self) -> int:
18         rand = randrange(0, self.s)
19         l, r = 0, len(self.enum)-1
20         e = self.enum
21         while l <= r:
22             m = (l + r) // 2
23             if e[m][0] <= rand < e[m][1]:
24                 return e[m][2]
25             elif rand >= e[m][1]:
26                 l = m+1
27             else:
28                 r = m-1
29         return e[m][0]
30 
31 from collections import Counter
32 obj = Solution([2, 3])
33 counter = Counter()
34 
35 for i in range(int(1e4)):
36     param_1 = obj.pickIndex()
37     counter[param_1] += 1
38 
39 print(counter)