~/problems / Binary search / Prefix sums + binary search

Random Pick with Weight

medium ~25 min

Implement a class WeightedPicker:

  • WeightedPicker(weights): weights is a non-empty list of non-negative integers, at least one of them positive.
  • pick_index(): return a random index i with probability weights[i] / sum(weights). An index with weight 0 must never be returned.

Example: with WeightedPicker([2, 5, 3]), pick_index() returns 0 about 20% of the time, 1 about 50% and 2 about 30%.

Constraints: up to 10^5 weights, each up to 10^5, and up to 10^5 calls to pick_index.

Use the global functions of Python's random module (random.random(), random.randint(...) and so on) so the tests can seed it. The tests check frequencies over many draws, with generous tolerance.

Each pick must be O(log n). Scanning the weights linearly on every pick, or calling random.choices(..., weights=...) each time, is O(n) per pick and fails the large test.

Show hint

do the O(n) work once, in the constructor: lay the weights end to end on a number line, so each index owns a stretch as long as its weight. A pick is then one random point on that line and a question about which stretch it falls in.

Topic: Prefix sums + binary search. Weighted sampling: bisect into cumulative weights.

Read the visual guide
0:00
Ctrl ' run · Ctrl ↵ submit
esc