Problem Description

You are given an integer array nums that may contain duplicate values. Your task is to implement a class that can randomly select and return an index where a given target number appears in the array.

The class should support two operations:

1. Constructor Solution(int[] nums): Initialize the object with the input array nums.

2. Pick Method int pick(int target): Given a target value, return a random index i where nums[i] == target. If the target appears multiple times in the array, each valid index should have an equal probability of being selected.

Key requirements:

• The target number is guaranteed to exist in the array
• When multiple indices contain the target value, the selection must be uniformly random (each valid index has equal chance of being picked)
• The array may contain duplicates

For example, if nums = [1, 2, 3, 3, 3] and you call pick(3), the method should randomly return one of the indices 2, 3, or 4, with each having a 1/3 probability of being selected.

The solution uses Reservoir Sampling, a technique for randomly selecting k items from a stream of unknown size. In this case, we're selecting 1 item (k=1) from all indices that match the target. As we iterate through the array:

• We maintain a count n of how many times we've seen the target
• For each occurrence, we randomly decide with probability 1/n whether to update our answer to the current index
• This ensures each valid index has equal probability of being the final result

Intuition

The challenge here is that we need to randomly select from indices we don't know in advance. We could scan the array first to find all indices where the target appears, store them, and then randomly pick one - but this would require extra space proportional to the number of occurrences.

Instead, we can use a clever approach that processes the array in a single pass while maintaining equal probability for each valid index. Think of it this way: imagine you're going through the array and you encounter the target values one by one.

• When you see the first occurrence, you have no choice - it becomes your current answer with probability 1/1 = 100%.
• When you see the second occurrence, you want both indices to have equal chance. So you keep the new index with probability 1/2 and keep the old one with probability 1/2.
• When you see the third occurrence, you want all three to have equal chance (1/3 each). If you pick the new one with probability 1/3, and keep the old answer with probability 2/3, then each of the three indices ends up with 1/3 probability overall.

This pattern continues: for the n-th occurrence, we select it with probability 1/n and keep the previous answer with probability (n-1)/n.

Mathematically, this works because:

• The first element has probability: 1 × 1/2 × 2/3 × 3/4 × ... × (n-1)/n = 1/n
• The second element has probability: 1/2 × 2/3 × 3/4 × ... × (n-1)/n = 1/n
• The k-th element has probability: 1/k × k/(k+1) × ... × (n-1)/n = 1/n

Each valid index ends up with exactly 1/n probability of being selected, achieving uniform randomness without needing to store all indices.

Solution Approach

The implementation uses the Reservoir Sampling algorithm (specifically for k=1) to achieve uniform random selection in a single pass through the array.

Data Structure:

• We only store the original array nums in the constructor
• No additional data structures are needed to track indices

Implementation Details:

1. Constructor: Simply stores the input array for later use.

   def __init__(self, nums: List[int]):
       self.nums = nums

2. Pick Method: Implements the reservoir sampling algorithm:

   def pick(self, target: int) -> int:
       n = ans = 0
       for i, v in enumerate(self.nums):
           if v == target:
               n += 1
               x = random.randint(1, n)
               if x == n:
                   ans = i
       return ans

Algorithm Walkthrough:

1. Initialize two variables:

   • n: Counter for how many times we've seen the target
   • ans: Stores the current selected index

2. Iterate through the array with index i and value v:

   • When we find a value equal to target:
     • Increment the counter n
     • Generate a random number x between 1 and n (inclusive)
     • If x == n, update ans to the current index i
     • This gives the current index a 1/n probability of being selected

3. Return the final value of ans

Why this works:

• The condition if x == n has probability 1/n of being true
• This means each new occurrence has exactly 1/n chance of replacing the previous answer
• Through mathematical induction, this ensures each valid index has equal probability 1/total_occurrences of being the final answer

Time and Space Complexity:

• Time: O(n) for each pick operation, where n is the length of the array
• Space: O(1) additional space (not counting the input array storage)

Example Walkthrough

Let's walk through the algorithm with nums = [1, 2, 3, 3, 3] and calling pick(3).

Initial State:

• n = 0 (count of target occurrences found)
• ans = 0 (current selected index)

Iteration Process:

Index 0, value = 1:

• 1 ≠ 3, skip

Index 1, value = 2:

• 2 ≠ 3, skip

Index 2, value = 3:

• Found target! n = 1
• Generate random number between 1 and 1: x = 1
• Since x == n (both are 1), update ans = 2
• Current probability: Index 2 selected with 100% chance

Index 3, value = 3:

• Found target! n = 2
• Generate random number between 1 and 2: let's say x = 2
• Since x == n (both are 2), update ans = 3
• Current probability distribution:
  • Index 2: had 100% × (1/2 chance of not being replaced) = 50%
  • Index 3: 50% chance of being selected

Index 4, value = 3:

• Found target! n = 3
• Generate random number between 1 and 3: let's say x = 1
• Since x ≠ n (1 ≠ 3), keep ans = 3
• Final probability distribution:
  • Index 2: 50% × (2/3 chance of surviving) = 1/3
  • Index 3: 50% × (2/3 chance of surviving) = 1/3
  • Index 4: 1/3 chance of being selected

Result: Returns ans = 3 in this run

Different Random Outcomes: If we ran this multiple times with different random values:

• Run 1: x=1 at n=1, x=1 at n=2, x=3 at n=3 → Returns index 4
• Run 2: x=1 at n=1, x=2 at n=2, x=2 at n=3 → Returns index 3
• Run 3: x=1 at n=1, x=1 at n=2, x=1 at n=3 → Returns index 2

Over many runs, each of the three indices (2, 3, 4) would be returned approximately 33.33% of the time, achieving uniform randomness.

Solution Implementation

Time and Space Complexity

Time Complexity: O(n) where n is the length of the input array nums.

• The __init__ method takes O(1) time as it only stores a reference to the input array.
• The pick method iterates through the entire array once using enumerate(), performing constant-time operations for each element (comparison, increment, random number generation, and conditional assignment). Therefore, the time complexity is O(n).

Space Complexity: O(1) for the pick method itself, O(n) for the overall solution.

• The __init__ method stores the input array, which requires O(n) space.
• The pick method uses only a constant amount of extra space with variables n, ans, i, v, and x, regardless of the input size.
• Note: This implementation uses reservoir sampling algorithm, which allows selecting a random index with equal probability without storing all target indices separately.

Common Pitfalls

1. Using 0-based Random Range Instead of 1-based

A common mistake is using random.randint(0, count-1) and checking if random_num == 0 instead of the 1-based approach.

Incorrect Implementation:

def pick(self, target: int) -> int:
    count = 0
    result_index = 0
  
    for index, value in enumerate(self.nums):
        if value == target:
            count += 1
            random_num = random.randint(0, count-1)  # Wrong range!
            if random_num == 0:  # This doesn't give 1/count probability
                result_index = index
  
    return result_index

Why it's wrong: With randint(0, count-1) when count=1, the range is just [0], so the probability is 100%. When count=2, the range is [0,1], giving 50% probability. This doesn't match the required 1/n probability pattern.

Correct Approach: Use random.randint(1, count) and check if random_num == count to ensure exactly 1/n probability.

2. Pre-computing All Indices in Constructor

Some might try to optimize by storing all indices for each value upfront:

Memory-Inefficient Approach:

def __init__(self, nums: List[int]):
    self.indices_map = {}
    for i, num in enumerate(nums):
        if num not in self.indices_map:
            self.indices_map[num] = []
        self.indices_map[num].append(i)

def pick(self, target: int) -> int:
    indices = self.indices_map[target]
    return random.choice(indices)

Problems:

• Uses O(n) additional space in worst case
• Defeats the purpose of reservoir sampling's space efficiency
• May cause memory issues with large arrays

3. Forgetting to Initialize Result Variable

Not initializing the result variable or using -1 as default can cause issues:

Problematic Code:

def pick(self, target: int) -> int:
    count = 0
    # result_index not initialized or initialized to -1
  
    for index, value in enumerate(self.nums):
        if value == target:
            count += 1
            if random.randint(1, count) == count:
                result_index = index  # NameError if not initialized!
  
    return result_index  # May return -1 or undefined

Solution: Always initialize result_index = 0 or to the first valid index found.

4. Using Wrong Probability Check

Some implementations incorrectly check if the random number equals 1 instead of count:

Wrong Probability Logic:

def pick(self, target: int) -> int:
    count = 0
    result_index = 0
  
    for index, value in enumerate(self.nums):
        if value == target:
            count += 1
            if random.randint(1, count) == 1:  # Wrong! This gives different probabilities
                result_index = index
  
    return result_index

Why it's wrong: This gives the first occurrence probability 1, second occurrence probability 1/2, third occurrence probability 1/3, etc. The final probabilities won't be uniform.

Correct Logic: Check if random.randint(1, count) == count to maintain the reservoir sampling property.