1from random import randrange
2from itertools import accumulate
3
4class 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
31from collections import Counter
32obj = Solution([2, 3])
33counter = Counter()
34
35for i in range(int(1e4)):
36 param_1 = obj.pickIndex()
37 counter[param_1] += 1
38
39print(counter)