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(vector<int>& nums) Initializes the object with the array nums.
  • int pick(int target) Picks a random index i from nums where nums[i] == target. If there are multiple valid i’s, then each index should have an equal probability of returning.

Examples

Example 1:

Input
["Solution", "pick", "pick", "pick"]
[[[1, 2, 3, 3, 3]], [3], [3], [3]]
Output
[null, 4, 4, 2]

Explanation
Solution solution = new Solution([1, 2, 3, 3, 3]);
solution.pick(3); // It should return either index 2, 3, or 4 randomly. Each index should have equal probability of returning.
solution.pick(3); // It should return either index 2, 3, or 4 randomly. Each index should have equal probability of returning.
solution.pick(3); // It should return either index 2, 3, or 4 randomly. Each index should have equal probability of returning.

Example 2:

Input
["Solution", "pick", "pick", "pick"]
[[[1, 2, 3, 3, 3]], [1], [3], [3]]
Output
[null, 0, 4, 2]

Explanation
solution.pick(1); // Should return 0. Since there's only one element with value 1 in the array.
solution.pick(3); // Should return either index 2, 3, or 4 randomly.
solution.pick(3); // Should return either index 2, 3, or 4 randomly.

Constraints

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

Thinking Process

  1. Preprocessing pays off: One-time O(n) cost amortized over many pick() calls
  • Identify required operations and their frequency (get/put/insert).
  • Combine data structures: hash map + list, heap + map, trie + DFS.
  • Amortized O(1) often needs lazy cleanup or doubly-linked lists.
Design pattern API hash + list compose data structures for operations

Common Approaches

Typical techniques for this pattern:

Approach Time Space Notes
Hash map + list (this problem) O(1) avg O(n) LRU cache pattern
Heap + hash map O(log n) O(n) LFU, time-based store
Trie (prefix tree) O(m) O(nm) Word search, autocomplete
Deque / circular buffer O(1) O(n) Queue with fixed capacity

Solution

Time Complexity: O(n) preprocessing, O(1) per pick() call
Space Complexity: O(n)

The key insight is to preprocess the array once during construction, storing all indices for each value in a hash map. This allows O(1) random selection per pick() call.

Solution: Hash Map Approach (Optimized for Multiple Picks)

from collections import defaultdict
import random


class Solution:
    def __init__(self, nums: list[int]):
        self.pos: dict[int, list[int]] = defaultdict(list)
        for i, x in enumerate(nums):
            self.pos[x].append(i)

    def pick(self, target: int) -> int:
        indices = self.pos[target]
        return indices[random.randrange(len(indices))]

Solution Explanation

Approach: Hash map + list (this problem)

Key idea: 1. Preprocessing pays off: One-time O(n) cost amortized over many pick() calls

How the code works:

  1. Preprocessing pays off: One-time O(n) cost amortized over many pick() calls
    • Identify required operations and their frequency (get/put/insert).
    • Combine data structures: hash map + list, heap + map, trie + DFS.
    • Amortized O(1) often needs lazy cleanup or doubly-linked lists.

      Algorithm Breakdown

Constructor: Build Index Map

from collections import defaultdict


def build_index_map(nums: list[int]) -> dict[int, list[int]]:
    pos: dict[int, list[int]] = defaultdict(list)
    for i, x in enumerate(nums):
        pos[x].append(i)
    return pos

Why:

  • Iterate through array once: O(n) time
  • For each value, store all its indices
  • Hash map provides O(1) average lookup time

Pick Function: Random Selection

import random


def pick_from_map(pos: dict[int, list[int]], target: int) -> int:
    indices = pos[target]
    return indices[random.randrange(len(indices))]

Why:

  • Lookup indices for target: O(1) average case
  • Random selection: rand() % size gives uniform distribution
  • Return selected index: O(1) total time

Why Reservoir Sampling Times Out

Solution 1: Reservoir Sampling (O(n) per pick)

import random


class Solution:
    def __init__(self, nums: list[int]):
        self.nums = nums

    def pick(self, target: int) -> int:
        rtn = -1
        cnt = 0
        for i, x in enumerate(self.nums):
            if x == target:
                cnt += 1
                if random.randrange(cnt) == 0:
                    rtn = i
        return rtn

Time Complexity: O(n) per pick() call
Space Complexity: O(1)

Why it times out:

  • Each pick() call scans entire array: O(n)
  • If pick() is called k times: O(k × n) total time
  • With k = 10^4 and n = 2×10^4: 200 million operations → TIMEOUT

When to use:

  • Single or few pick() calls
  • Memory-constrained environments
  • Streaming data (can’t preprocess)

Complexity Comparison

Approach Constructor pick() Total (k calls) Space Best For
Hash Map O(n) O(1) O(n + k) O(n) Multiple picks
Reservoir Sampling O(1) O(n) O(k × n) O(1) Single pick

Key Insight:

  • Hash map: Amortized cost per pick is O(1) after preprocessing
  • Reservoir sampling: Per-call cost is O(n), expensive for many calls

Implementation Details

Why Hash Map is Optimal

For multiple pick() calls:

  • Preprocessing: O(n) one-time cost
  • Each pick: O(1) lookup + O(1) random selection
  • Total for k picks: O(n + k) vs O(k × n) for reservoir sampling

Example:

n = 20,000, k = 10,000 picks

Hash Map:
  Constructor: 20,000 operations
  Picks: 10,000 operations
  Total: 30,000 operations ✓

Reservoir Sampling:
  Constructor: 0 operations
  Picks: 10,000 × 20,000 = 200,000,000 operations ✗ TIMEOUT

Random Number Generation

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
                if random.randrange(count) == 0:
                    result = i
        return result

Why this works:

  • rand() returns pseudo-random integer
  • % size maps to range [0, size-1]
  • Uniform distribution (assuming good RNG)

Note: For better randomness, consider:

import random


class Solution:
    def __init__(self, nums: list[int]):
        self.nums = nums

    def pick(self, target: int) -> int:
        indices = [i for i, x in enumerate(self.nums) if x == target]
        return indices[random.randrange(len(indices))]

Common Mistakes

  1. Single occurrence: pick() returns the only index
  2. All same values: All indices have equal probability
  3. Large array: Hash map handles efficiently

  4. Many pick calls: Hash map approach scales better

  5. Using reservoir sampling for many picks: Times out
  6. Not preprocessing: Rebuilding indices each time
  7. Wrong random selection: Not using modulo correctly
  8. Memory concerns: Hash map uses O(n) space, but worth it for speed
  9. Assuming single occurrence: Must handle multiple indices

Optimization Tips

  1. Preprocess in constructor: Build hash map once
  2. Use reference: auto& indices avoids copying
  3. Hash map over array: O(1) lookup vs O(n) scan
  4. Consider when to use each:
    • Many picks → Hash map
    • Single pick → Reservoir sampling

Real-World Applications

  1. Load Balancing: Randomly select server from pool
  2. Sampling: Randomly sample data points
  3. Game Development: Random item/spawn selection
  4. Testing: Random test case generation

Pattern Recognition

This problem demonstrates the “Preprocessing for Fast Lookup” pattern:

1. Identify repeated operations (pick() called many times)
2. Preprocess data once (build hash map in constructor)
3. Optimize each operation (O(1) lookup instead of O(n) scan)
4. Trade space for time (O(n) space for O(1) time)

Similar problems:

  • Design data structures with fast operations
  • Caching frequently accessed data
  • Precomputation for repeated queries

Why Hash Map Solution is Optimized

Time Analysis:

  • Constructor: O(n) - one-time cost
  • pick(): O(1) average case - hash lookup + random selection
  • Total for k picks: O(n + k) - linear in total operations

Space Analysis:

  • O(n) - stores all indices (at most n indices total)

Comparison:

  • Reservoir sampling: O(k × n) time for k picks
  • Hash map: O(n + k) time for k picks
  • Improvement: From O(k × n) to O(n + k) = k times faster for k picks

Step-by-Step Trace: nums = [1, 2, 3, 3, 3], pick(3) called 3 times

Constructor:
  i=0: nums[0]=1 → pos[1] = [0]
  i=1: nums[1]=2 → pos[2] = [1]
  i=2: nums[2]=3 → pos[3] = [2]
  i=3: nums[3]=3 → pos[3] = [2, 3]
  i=4: nums[4]=3 → pos[3] = [2, 3, 4]

pick(3) call 1:
  indices = pos[3] = [2, 3, 4]
  random = rand() % 3 = 1
  return indices[1] = 3

pick(3) call 2:
  indices = pos[3] = [2, 3, 4]
  random = rand() % 3 = 0
  return indices[0] = 2

pick(3) call 3:
  indices = pos[3] = [2, 3, 4]
  random = rand() % 3 = 2
  return indices[2] = 4

Mathematical Proof: Uniform Distribution

Hash Map Approach:

  • Each index in indices array has equal probability
  • rand() % size gives uniform distribution over [0, size-1]
  • Therefore, each index has probability 1/size ✓

Reservoir Sampling:

  • For k occurrences, each index i has probability:
    • Selected at position i: 1/i
    • Not replaced by later indices: (i/(i+1)) × ((i+1)/(i+2)) × ... × ((k-1)/k)
    • Product = 1/k ✓

Both approaches guarantee uniform distribution!

Key Takeaways

  1. Preprocessing pays off: One-time O(n) cost amortized over many pick() calls
  2. Hash map lookup: O(1) average case for finding indices
  3. Random selection: Use rand() % size for uniform distribution
  4. Space-time tradeoff: Use O(n) space to achieve O(1) time per pick

References

Template Reference