Random Pick Index — Reservoir Sampling for Uniform Random Selection

Problem Statement#

Given an integer array nums with possible duplicates, randomly output the index of a given target number. You can assume that the given target number must exist in the array.

Implement the Solution class:

• Solution(int[] nums) initialises the object with the array.
• int pick(int target) picks a random index i from nums where nums[i] == target. If there are multiple valid indices, each index should have an equal probability of returning.

Constraints:

• 1 <= nums.length <= 2 * 10^4
• -2^31 <= nums[i] <= 2^31 - 1
• target is an integer in the range [-2^31, 2^31 - 1].
• At most 10^4 calls will be made to pick.

Examples:

Input: ["Solution","pick","pick","pick"]
       [[[1,2,3,3,3]],[3],[1],[3]]
Output: [null,4,0,2]
Explanation:
  - pick(3): could return 2, 3, or 4 with equal probability (1/3 each)
  - pick(1): must return 0 (only one index)
  - pick(3): could return 2, 3, or 4 with equal probability
 
Input: nums = [1,2,3,3,3], pick(3) called many times.
Expected: indices 2, 3, 4 each returned roughly 1/3 of the time.

Why This Problem Matters#

Reservoir sampling is a fundamental algorithm for uniform sampling from a data stream of unknown size. It appears in Google and Facebook interviews because it requires probabilistic reasoning combined with algorithm design — skills that are tested in roles ranging from data engineering to distributed systems.

The problem has two natural approaches:

1. Pre-process: Store target → [list of indices] in a hash map during construction. For pick(target), randomly select from the list. O(n) construction, O(1) pick, O(n) space.

2. Reservoir sampling: No pre-processing. During pick, iterate through the array and sample. O(1) construction (just store the array), O(n) pick, O(1) extra space.

The reservoir sampling approach is the interview answer because it demonstrates the ability to design for constraints that aren't obvious: what if the array is too large to fit in memory? What if the array changes between calls (a stream)? Reservoir sampling handles both scenarios because it works on any stream without buffering.

Google uses this problem to evaluate candidates for streaming data infrastructure roles. Facebook (Meta) uses it in systems design discussions about fair A/B test assignment. At Amazon, it connects to the broader topic of probabilistic data structures.

The Core Insight#

Reservoir Sampling Proof:

We want each index i where nums[i] == target to be selected with probability 1/k where k is the total count of target in the array. We don't know k in advance.

Algorithm: iterate through the array, maintaining count = 0 and result = -1. When nums[i] == target, increment count and replace result with i with probability 1/count.

Why does this work?

By induction. After seeing the first occurrence: result = first_index with probability 1 (1/1). Correct since k=1 at this point.

After seeing the second occurrence: The new index replaces with probability 1/2. The first index survives with probability 1 - 1/2 = 1/2. Both have probability 1/2. Correct since k=2.

After seeing the k-th occurrence: The k-th index is selected with probability 1/k. Each earlier index i (1 ≤ i ≤ k-1) was selected after i occurrences with probability 1/i, and survived steps i+1, i+2, ..., k (not replaced) with probability:

(1 - 1/(i+1)) * (1 - 1/(i+2)) * ... * (1 - 1/k)
= i/(i+1) * (i+1)/(i+2) * ... * (k-1)/k
= i/k

So probability that index i is selected = (1/i) × (i/k) = 1/k. All k indices have probability 1/k. The sampling is uniform.

Visual Dry Run#

nums = [1, 2, 3, 3, 3], pick(3):

After all iterations: result is 2, 3, or 4, each with probability 1/3.

Solution (Optimal)#

import random
 
class Solution:
    def __init__(self, nums: list[int]):
        self.nums = nums
 
    def pick(self, target: int) -> int:
        count = 0
        result = -1
 
        for i, x in enumerate(self.nums):
            if x == target:
                count += 1
                # Replace result with probability 1/count
                if random.randint(1, count) == 1:
                    result = i
 
        return result

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

Alternative — Pre-processing with HashMap:

from collections import defaultdict
import random
 
class Solution:
    def __init__(self, nums: list[int]):
        self.index_map = defaultdict(list)
        for i, x in enumerate(nums):
            self.index_map[x].append(i)
 
    def pick(self, target: int) -> int:
        indices = self.index_map[target]
        return random.choice(indices)

Complexity Analysis:

Reservoir sampling:

• Constructor: O(1) — just store the array reference.
• pick: O(n) — scan the entire array.
• Space: O(1) extra (plus O(n) for the stored array).

Pre-processing with hash map:

• Constructor: O(n) — build the index map.
• pick: O(1) — direct lookup and random choice.
• Space: O(n) — store all indices.

Trade-off: If pick is called rarely, reservoir sampling wastes less time in the constructor. If pick is called frequently, the hash map approach pays off.

Common Mistakes#

• Using random.randint(0, count-1) == 0 in Python: This is equivalent to randint(1, count) == 1 — both give probability 1/count. Just be consistent and understand which form you're using.
• Initialising result = 0 instead of result = -1: If the problem guarantees target exists, result will always be set. But initialising to -1 makes it clear what "not found" means.
• Confusing random.randint behaviour across languages: In Python, random.randint(1, count) is inclusive on both ends (gives values 1, 2, ..., count). In Java, rand.nextInt(count) gives values 0, 1, ..., count-1. In JavaScript, Math.floor(Math.random() * count) gives 0 to count-1. All give probability 1/count of selecting 0 (or 1).
• Not storing the array in the constructor: The reservoir sampling approach needs access to the array during pick. Don't deep copy if the array is large — store a reference.
• Testing uniformity: When verifying your implementation, call pick thousands of times and check that each valid index appears approximately equally. Don't just test correctness for one call.

Follow-up Questions#

1. How would you extend reservoir sampling to select k items (k > 1) uniformly at random from a stream? (Classic k-reservoir algorithm: fill the first k elements, then for each subsequent element i, replace a random element among the first i with probability k/i.)
2. What if the array is too large to fit in memory — how does reservoir sampling help? (Stream the array in chunks; reservoir sampling works on any stream without needing to buffer all elements.)
3. What if the array can change (elements added or removed)? (Reservoir sampling doesn't help here; use a dynamic data structure like a balanced BST or a Fenwick tree.)
4. Can you make pick O(log n) instead of O(n)? (Yes: preprocess into a hash map of target → sorted index list, then use binary search — but pick is already O(1) with the hash map approach.)
5. How does reservoir sampling relate to Fisher-Yates shuffle? (Both achieve uniform permutation/selection; Fisher-Yates is a one-pass shuffle of a known-size array.)
6. What is the expected number of times result gets updated during pick(target) when there are k occurrences? (Harmonic number H(k) = 1 + 1/2 + 1/3 + ... + 1/k ≈ ln(k) updates on average.)

Key Takeaways#

• LC 398 Random Pick Index uses reservoir sampling to draw a uniform index without storing all matches.
• For each i where nums[i] == target, replace the running answer with probability 1/count where count is the running match count.
• Equivalent: pick a random integer in [0, count) — if zero, set result = i.
• Each candidate ends up selected with probability exactly 1/k (where k is total matches), proven by a simple induction.
• O(1) extra space — never materializes the full list of matching indices.
• Each pick(target) is O(n); preprocessing (__init__) is O(1) — ideal when picks are rare and memory is tight.
• Reservoir sampling generalizes to LC 382 (linked list random node), random log line picking, and streaming statistics.

Related Problems#

• [LC 398] Random Pick Index — this exact problem.
• [LC 382] Linked List Random Node — reservoir sampling on a linked list (same algorithm, different data structure).
• [LC 384] Shuffle an Array — Fisher-Yates shuffle, related random permutation concept.
• [LC 528] Random Pick with Weight — weighted random selection using prefix sums + binary search.
• [LC 519] Random Flip Matrix — uniform random selection from a dynamic set of positions.
• [LC 710] Random Pick with Blacklist — random selection with excluded values.