We use cookies (including Google cookies) to personalize ads and analyze traffic. By continuing to use our site, you accept our Privacy Policy.

SimplyLeet
TagsCompaniesDisclaimersPrivacyTermsContactAbout Us

Random Pick Index

Number: 398

Difficulty: Medium

Paid? No

Companies: Meta, Amazon, TikTok

Problem Description

Initialize the object with an integer array nums (which may contain duplicates) and implement the method pick(target) that randomly returns an index i from nums such that nums[i] == target. Each valid index for a target must have equal probability of being returned.

Key Insights

  • Reservoir sampling is used to select a random index without using extra space for storing indices.
  • During a single pass over the array, you can decide probabilistically whether to choose the current valid index.
  • Consistent probability handling ensures each target occurrence is equally likely.
  • This approach avoids preprocessing, keeping space usage minimal.

Space and Time Complexity

Time Complexity: O(n) per pick call, where n is the length of the nums array. Space Complexity: O(1) extra space (apart from storing the input array).

Solution

We leverage the reservoir sampling algorithm for each pick call. When the pick method is called with a target value:

  1. Initialize a counter (found) to track the number of target occurrences.
  2. Iterate over the nums array.
  3. Each time the target is found, increment the counter.
  4. With probability 1/found, replace the result with the current index.
  5. Return the final selected index.

This method ensures that when multiple valid indices exist, each one is chosen with the same probability.

Code Solutions

import random

class Solution:
    def __init__(self, nums: List[int]):
        # Store the input array for later pick calls
        self.nums = nums

    def pick(self, target: int) -> int:
        # Initialize count and result index
        count = 0
        result_index = -1
        # Iterate over every index and value in the array
        for index, value in enumerate(self.nums):
            if value == target:
                count += 1  # Increment count for each occurrence of target
                # Randomly decide whether to pick this index with probability 1/count
                if random.randint(1, count) == 1:
                    result_index = index
        return result_index
← Back to All Questions