Code Implementation

import random

class Solution:
    def __init__(self, nums):
        self.nums = nums

    def pick(self, target):
        count = 0
        res = -1
        for i, num in enumerate(self.nums):
            if num == target:
                count += 1
                if random.randint(1, count) == 1:
                    res = i
        return res
      

#include <vector>
#include <cstdlib>

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

import java.util.*;

class Solution {
    private int[] nums;
    private Random rand;

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

var Solution = function(nums) {
    this.nums = nums;
};

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

Problem Description

You are given an integer array nums that may contain duplicates. Implement a class with a method pick(target) that returns a random index where the value of nums at that index is equal to target. Each valid index should have an equal probability of being returned.

Key constraints:

• There may be multiple indices with the value target.
• Each call to pick(target) must return one of these indices, chosen uniformly at random.
• The array can be large, and the pick method may be called many times.
• Do not use extra space proportional to the size of nums for preprocessing.

Thought Process

The first idea that might come to mind is to store all indices of target in a separate list, and then randomly pick one. But this requires extra space proportional to the number of occurrences of target, and if the array is very large or the method is called frequently, this can be inefficient.

Instead, we need a way to pick a random index "on the fly" without creating a list of all valid indices. This is a classic case for an algorithm called Reservoir Sampling, which allows us to select one random item from a stream (or array) where the number of qualifying items is not known in advance.

The challenge is to ensure that every index with value target has an equal chance of being selected, even as we scan through the array only once.

Solution Approach

We use the Reservoir Sampling algorithm for this problem. Here’s how it works step-by-step:

1. Initialize a variable to count how many times we've seen target so far (count), and another to store the current chosen index (res).
2. Iterate through the array. For each index i:
   • If nums[i] is not equal to target, skip it.
   • If nums[i] equals target, increment count.
   • Generate a random integer between 1 and count (inclusive). If the random number is 1 (or 0 depending on language), update res to i.
3. After finishing the loop, res holds the index of a randomly selected occurrence of target.

This approach ensures that each valid index is chosen with equal probability, and it only requires constant extra space regardless of the input size.

Why Reservoir Sampling? Because it allows us to pick a random item from a stream (or array) where the total number of qualifying items is not known beforehand, and it does so using only constant space.

Example Walkthrough

Suppose nums = [1, 2, 3, 3, 3] and we call pick(3).

• We scan through the array:
• At index 0: nums[0]=1 (not 3), skip.
• At index 1: nums[1]=2 (not 3), skip.
• At index 2: nums[2]=3 (first occurrence of 3): count=1. Random number between 1 and 1 is always 1, so res=2.
• At index 3: nums[3]=3 (second occurrence): count=2. Random number between 1 and 2. If 1, res=3; else, keep previous.
• At index 4: nums[4]=3 (third occurrence): count=3. Random number between 1 and 3. If 1, res=4; else, keep previous.

At the end, res will be one of 2, 3, or 4, each with probability 1/3. If you run pick(3) many times, you'll see each index is chosen about equally often.

Time and Space Complexity

• Brute-force approach: For each pick(target) call, collect all indices where nums[i] == target into a list (O(n) time and O(n) space), then pick one at random (O(1) time).
• Optimized (Reservoir Sampling) approach: For each pick(target) call, scan the array once (O(n) time), but use only O(1) extra space.

Since the problem restricts extra space, the reservoir sampling solution is optimal: O(n) time per pick, O(1) space.

Summary

The Random Pick Index problem is elegantly solved using reservoir sampling, allowing us to select a random index of a target value with equal probability and without using extra space proportional to the input size. The key insight is that, by updating the chosen index with decreasing probability as we see more occurrences, we ensure fairness and efficiency. This approach is both simple and powerful, and is widely applicable in streaming and sampling scenarios.