Random Pick Index

medium math reservoir sampling randomized

Problem

Design a class that, given an array nums and a query target, returns a uniformly random index i where nums[i] == target. Use reservoir sampling so memory stays O(1) per query.

Inputnums = [1, 2, 3, 3, 3], target = 3

Outputone of {2, 3, 4} each with prob 1/3

Stream over nums; the k-th match replaces the reservoir with probability 1/k.

nums (comma-separated)

target

import random
class Solution:
    def __init__(self, nums):
        self.nums = nums
    def pick(self, target):
        chosen, count = -1, 0
        for i, v in enumerate(self.nums):
            if v == target:
                count += 1
                if random.randrange(count) == 0:
                    chosen = i
        return chosen

class Solution {
  constructor(nums) { this.nums = nums; }
  pick(target) {
    let chosen = -1, count = 0;
    for (let i = 0; i < this.nums.length; i++) {
      if (this.nums[i] === target) {
        count++;
        if (Math.floor(Math.random() * count) === 0) chosen = i;
      }
    }
    return chosen;
  }
}

class Solution {
    int[] nums; Random rng = new Random();
    public Solution(int[] nums) { this.nums = nums; }
    public int pick(int target) {
        int chosen = -1, count = 0;
        for (int i = 0; i < nums.length; i++) {
            if (nums[i] == target) {
                count++;
                if (rng.nextInt(count) == 0) chosen = i;
            }
        }
        return chosen;
    }
}

class Solution {
public:
    vector<int> nums;
    Solution(vector<int>& n) : nums(n) {}
    int pick(int target) {
        int chosen = -1, count = 0;
        for (int i = 0; i < (int)nums.size(); i++) {
            if (nums[i] == target) {
                count++;
                if (rand() % count == 0) chosen = i;
            }
        }
        return chosen;
    }
};

Explanation

We want to return a random index where nums[i] == target, with each matching index equally likely. The neat part: we do it in one pass without ever storing the list of matches, using a trick called reservoir sampling.

We keep a running count of how many matches we have seen and a single chosen index. Each time we hit the k-th match, we replace chosen with the current index with probability 1/k — that is what random.randrange(count) == 0 does.

Why is this fair? The last match (the k-th) is kept with probability 1/k. The earlier ones survive because each later replacement only happens sometimes, and the probabilities multiply out so that every match ends up with the same chance, 1/k overall.

Example: nums = [1,2,3,3,3], target = 3. The matches are at indices 2, 3, 4. Index 2 is taken first; at the 2nd match it is kept with prob 1/2, at the 3rd with prob 2/3. Each of indices 2, 3, 4 ends up chosen with probability exactly 1/3.

This uses only O(1) extra memory since we never build a list of positions.

Time: O(n) per pick Space: O(1) beyond input