Implement a class that picks an index at random, with probability proportional to its weight:
WeightedPicker(weights, rng):weightsis a list of positive ints;rngis arandom.Random. Use it for every random choice (e.g.rng.randrange(k)orrng.random()), never the globalrandommodule orrandom.choices.pick() -> int: return indexiwith probabilityweights[i] / sum(weights).
p = WeightedPicker([1, 3], random.Random(0))
p.pick() # 0 with probability 1/4, 1 with probability 3/4
p = WeightedPicker([2, 5, 3], random.Random(0))
p.pick() # 0, 1 or 2 with probabilities 0.2, 0.5, 0.3
Picture the weights laid end to end on a number line: index 0 owns [0, 2), index 1 owns [2, 7), index 2 owns [7, 10). Draw a random point r in [0, total) and return the index whose segment contains it. Precompute the prefix sums once in the constructor, then find the segment with a binary search so each pick is O(log n).
Constraints: 1 <= len(weights) <= 10^5, 1 <= weights[i] <= 10^5, up to 10^5 calls to pick.
Show hint
With prefix sums [2, 7, 10] and an integer r in [0, 10), the answer is bisect.bisect_right(prefix, r).