~/problems / Probability / Randomized algorithms

Random pick avoiding a blacklist

medium ~30 min

Pick uniformly random integers from [0, n) while avoiding a set of forbidden values.

  • Solution(n, blacklist, rng): blacklist holds distinct values in [0, n); rng is a random.Random. Use it for every random choice (e.g. rng.randrange(k)), never the global random module.
  • pick() -> int: return a value in [0, n) that is not blacklisted, each allowed value equally likely.

n can be up to 10^9, so you can't list the allowed values. And retrying until you miss the blacklist ("rejection sampling") is hopeless when almost everything is blacklisted. Aim for O(len(blacklist)) setup and O(1) per pick, calling the rng once.

s = Solution(7, [2, 3, 5], random.Random(0))
s.pick()   # one of 0, 1, 4, 6, each with probability 1/4

Constraints: 1 ≤ n ≤ 10^9, 0 ≤ len(blacklist) ≤ min(10^5, n - 1), up to 2·10^4 calls to pick.

Show hint

There are exactly m = n - len(blacklist) allowed values, so draw a number in [0, m). Most such numbers are already allowed; the blacklisted ones in that range can each be redirected, via a dict built once, to an allowed value in [m, n), and there are exactly enough of those.

Topic: Randomized algorithms. Reservoir sampling, Fisher-Yates, rejection sampling.

0:00
Ctrl ' run · Ctrl ↵ submit
esc