398. Random Pick Index

MediumReservoir SamplingHash TableMathRandomized

Problem Description

The problem presents a class called Solution that is initialized with an array of integers nums that may contain duplicates. The core functionality we need to implement is the pick method, which should return the index of a randomly selected occurrence of a target number provided as input. Importantly, each possible index that matches the target should have an equal probability of being chosen. This means that if the target number appears multiple times in the array, our method should not show any bias towards any specific index.

Intuition

The intuition behind the solution approach to this problem lies in ensuring each index with the target value has an equal chance of being selected. This is a direct application of a well-known algorithm called Reservoir Sampling. Reservoir Sampling is useful when we have a large or unknown size of input and we need to select a random element (or a subset of elements) from it with equal probability.

In this case, we're iterating through the entire array, and each time we find the target number, we treat it as a potential candidate to be our random pick. As we encounter each target, we increase a count n which keeps track of how many times we've seen the target number so far. We then generate a random number 'x' between 1 and n. If 'x' is equal to n, which statistically will be 1 in n times, we set the ans variable to the index i. Through this process, every index has an equal chance (1/n) of being chosen since the choice is re-evaluated each time we see the target number. By the end of the array, this ensures that each index where nums[i] == target has been chosen with equal probability.

Learn more about Math patterns.

Solution Approach

The solution employs a clever implementation of Reservoir Sampling to ensure that each potential index for the target value is chosen with equal probability. Let's step through the implementation provided in the Reference Solution Approach:

We start with the __init__ method of the Solution class, which simply stores the nums array as an instance variable for later use.
The pick method is where the logic for Reservoir Sampling comes into play. Here, we initialize two variables, n, which will track the number of times we've encountered the target value, and ans, which will hold our currently chosen index for the target.
We then loop through each index-value pair in nums using the enumerate function.
Inside the loop, we check if the current value v is equal to the target value. If it's not, we continue looping without doing anything further.
However, if v is the target, we increment our count n. This count is crucial because it influences the probability of selecting the current index.
Next, we generate a random number x between 1 and n using random.randint(1, n). This random number decides whether we should update our current answer ans to the current index i.
The condition if x == n: is the key to ensuring that each index has an equal chance of being chosen. This condition will be true exactly once in n occurrences on average, which aligns with the probability we want. When this condition is true, we set ans to i.
After the loop completes, ans will hold the index of one of the occurrences of the target value, chosen at random. We return ans.

In summary, this algorithm effectively iterates through the list of numbers once (O(n) time complexity), using a constant amount of space (O(1) space complexity, not counting the input array), handling the potential of an unknown number of duplicates in a way that each target index has an equal likelihood of being selected.

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's illustrate the solution approach with a small example. Suppose our nums array is [4, 3, 3, 2, 3], and we want to pick an index at random where the target value 3 is located.

We initialize the Solution class with nums.
We call the pick method with the target value 3.
The method begins to loop over nums. As it finds 3 at index 1, we increase n to 1 and generate a random number x. Since n is 1, x must also be 1. We set ans to the index of our target, which is 1.
The loop proceeds and finds another 3 at index 2. It increases n to 2 and generates a random number x between 1 and 2. If x turns out to be 1, we keep ans as it is (1), but if x is 2, we update ans to 2.
Another 3 appears at index 4. This time, n becomes 3, and we repeat the process. A random number x is chosen between 1 and 3. We have one-third of a chance to update ans to 4 if x is 3.

Assuming our random numbers for each step are 1, 2, and 1 respectively, here is how the process unfolds:

First occurrence (index 1): n is 1, x is 1, so ans is set to 1.
Second occurrence (index 2): n is 2, x is 2, so ans is updated to 2.
Third occurrence (index 4): n is 3, x is 1, no change to ans.

Thus, the final answer returned by our pick method in this example could be 2, which is one of the indexes where the target value 3 appears. The above iterations demonstrate that each index (1, 2, and 4) had an equal chance of being chosen. The probability for each was 1/3 by the end of the process.

Solution Implementation

1import random
2from typing import List
3
4class Solution:
5    def __init__(self, nums: List[int]):
6        """Initialize the Solution object with a list of numbers."""
7        self.nums = nums
8
9    def pick(self, target: int) -> int:
10        """
11        Pick a random index from the list of numbers where the number at that index equals the target.
12        This uses Reservoir Sampling to ensure that each such index has an equal probability of being chosen.
13        """
14        count = 0  # Counter for the occurrence of the target
15        chosen_index = None  # Variable to store the randomly chosen index
16
17        # Enumerate over the list of numbers
18        for index, value in enumerate(self.nums):
19            # Check if the current value matches the target
20            if value == target:
21                count += 1  # Increment the counter for each occurrence
22                # Generate a random number between 1 and the current count (both inclusive)
23                random_number = random.randint(1, count)
24                # If the generated number equals the current count, update the chosen index
25                if random_number == count:
26                    chosen_index = index
27        # Return the selected index which corresponds to the target in the list
28        return chosen_index
29
30
31# An example of how the Solution class might be instantiated and used:
32# solution_instance = Solution(nums)
33# random_index = solution_instance.pick(target)
34

1import java.util.Random;
2
3class Solution {
4    private int[] nums; // This array holds the original array of numbers.
5    private Random random = new Random(); // Random object to generate random numbers.
6
7    // Constructor that receives an array of numbers.
8    public Solution(int[] nums) {
9        this.nums = nums; // Initialize the nums array with the given input array.
10    }
11
12    // Method to pick a random index where the target value is found in the nums array.
13    public int pick(int target) {
14        int count = 0; // Counter to track how many times we've seen the target so far.
15        int result = 0; // Variable to keep the result index.
16
17        // Iterating over the array to find target.
18        for (int i = 0; i < nums.length; ++i) {
19            if (nums[i] == target) { // Check if current element is the target.
20                count++; // Increment the count since we have found the target.
21                // Generate a random number between 1 and the number of times target has been seen inclusively.
22                int randomNumber = 1 + random.nextInt(count);
23
24                // If the random number equals to the count (probability 1/n),
25                // set the result to current index i.
26                if (randomNumber == count) {
27                    result = i;
28                }
29            }
30        }
31        return result; // Return the index of target chosen uniformly at random.
32    }
33}
34
35/**
36 * The following is how to use the Solution class:
37 * Solution solution = new Solution(nums);
38 * int index = solution.pick(target);
39 */
40

1#include <vector>
2#include <cstdlib> // For rand()
3
4class Solution {
5private:
6    vector<int> elements; // Renamed 'nums' to 'elements' for clarity
7
8public:
9    // Constructor which initializes the 'elements' vector
10    Solution(vector<int>& nums) : elements(nums) {}
11
12    // Function to pick a random index for a given target
13    // This implements Reservoir Sampling
14    int pick(int target) {
15        int count = 0;      // Store the number of occurrences of 'target'
16        int chosenIndex = -1; // Store the randomly chosen index of 'target'
17
18        // Iterate through the elements to find 'target'
19        for (int i = 0; i < elements.size(); ++i) {
20            if (elements[i] == target) {
21                count++; // Increment count for each occurrence
22                // With probability 1/count, choose this index
23                if ((rand() % count) == 0) {
24                    chosenIndex = i;
25                }
26            }
27        }
28        // Return the randomly chosen index of 'target'
29        return chosenIndex;
30    }
31};
32
33/**
34 * Usage:
35 * vector<int> numbers = {1, 2, 3, 3, 3};
36 * Solution* solution = new Solution(numbers);
37 * int randomIndexForTarget = solution->pick(3);
38 * // randomIndexForTarget will have the index of '3' chosen uniformly at random
39 * delete solution; // Don't forget to free the allocated memory!
40 */
41

1class Solution {
2    // Array to store the elements
3    elements: number[];
4
5    constructor(nums: number[]) {
6        this.elements = nums;
7    }
8
9    pick(target: number): number {
10        let count: number = 0;       // Store the number of occurrences of 'target'
11        let chosenIndex: number = -1; // Store the randomly chosen index of 'target'
12
13        // Iterate through the elements to find occurrences of 'target'
14        for (let i = 0; i < this.elements.length; i++) {
15            if (this.elements[i] === target) {
16                count++; // Increment count for each occurrence
17                // With probability 1/count, choose this index
18                if (Math.floor(Math.random() * count) === 0) {
19                    chosenIndex = i;
20                }
21            }
22        }
23
24        // Return the randomly chosen index of 'target'
25        return chosenIndex;
26    }
27}
28
29/**
30* Your Solution object will be instantiated and called as such:
31* var obj = new Solution(nums)
32* var param_1 = obj.pick(target)
33*/
34

Time and Space Complexity

Time Complexity

The time complexity of the pick method is O(N), where N is the total number of elements in nums. This is because, in the worst case, we have to iterate over all the elements in nums to find all occurrences of target and decide whether to pick each occurrence or not.

During each iteration, we perform the following operations:

Compare the current value v with target.
If they match, we increment n.
Generate a random number x with random.randint(1, n), which has constant time complexity O(1).
Compare x to n and possibly update ans.

These operations are within the single pass loop through nums, hence maintaining the overall time complexity of O(N).

Space Complexity

The space complexity of the pick method is O(1). The additional space required for the method execution does not depend on the size of the input array but only on a fixed set of variables (n, ans, i, and v), which use a constant amount of space.

The class Solution itself has space complexity O(N), where N is the number of elements in nums, since it stores the entire list of numbers.

However, when analyzing the space complexity of the pick method, we consider only the extra space used by the method excluding the space used to store the input, which in this case remains constant.

Learn more about how to find time and space complexity quickly using problem constraints.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Question 1 out of 10

Which of the tree traversal order can be used to obtain elements in a binary search tree in sorted order?

Pre-order traversal

Post-order traversal

In-order traversal

Level-order traversal

I don’t know