~/problems / Probability / Randomized algorithms

Random index of a target value

easy ~15 min

You're given a list of ints that may contain duplicates. Implement:

  • Solution(nums, rng): rng is a random.Random; use it for every random choice (e.g. rng.randrange(k)), never the global random module.
  • pick(target) -> int: return an index i with nums[i] == target, chosen uniformly among all such indices. The target is guaranteed to be in the list.
s = Solution([3, 1, 3, 2, 3], random.Random(0))
s.pick(3)   # 0, 2 or 4, each with probability 1/3
s.pick(1)   # always 1

Constraints: 1 ≤ len(nums) ≤ 2·10^4, up to 10^4 calls to pick.

Follow-up: can you do it with O(1) extra memory, at the cost of O(n) per pick?

Show hint

Doing the work once in the constructor makes each pick cheap: remember where every value occurs, then choose uniformly among those positions. For the O(1)-memory follow-up, scan the list on each call and let each new match replace your current choice with a probability that keeps all matches seen so far equally likely.

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

0:00
Ctrl ' run · Ctrl ↵ submit
esc