Random Pick Index
MedPrompt
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)Initializes the object with the arraynums.int pick(int target)Picks a random indexifromnumswherenums[i] == target. If there are multiple valid i's, then each index should have an equal probability of returning.
Example 1:
Input
["Solution", "pick", "pick", "pick"]
[[[1, 2, 3, 3, 3]], [3], [1], [3]]
Output
[null, 4, 0, 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(1); // It should return 0. Since in the array only nums[0] is equal to 1.
solution.pick(3); // It should return either index 2, 3, or 4 randomly. Each index should have equal probability of returning.
Constraints:
1 <= nums.length <= 2 * 104-231 <= nums[i] <= 231 - 1targetis an integer fromnums.- At most
104calls will be made topick.
Approaches
3 approaches with complexity analysis and trade-offs.
This is a straightforward brute-force approach. For every call to pick(target), we iterate through the entire input array. We collect all the indices where the element equals the target into a temporary list. Finally, we randomly select one index from this list and return it. The constructor simply stores a reference to the input array.
Algorithm
- In the
pick(target)method, initialize an empty list to store the indices of the target element. - Iterate through the entire
numsarray from the beginning to the end. - For each element, check if it is equal to the
target. - If
nums[i] == target, add the current indexito the list. - After the iteration is complete, the list will contain all indices where the
targetappears. - Generate a random number between 0 (inclusive) and the size of the list (exclusive).
- Return the element from the list at the randomly generated index.
Walkthrough
The Solution class stores the input array. The pick method performs a linear scan of this array. It uses an auxiliary ArrayList to keep track of all indices i where nums[i] matches the target. Once the entire array has been scanned, it means we have found all possible indices. Then, a random index is chosen from this list of indices. Since every valid index is in the list, and we pick a random element from the list, each index has an equal probability of being chosen.
import java.util.ArrayList;import java.util.List;import java.util.Random; class Solution { private int[] nums; private Random rand; public Solution(int[] nums) { this.nums = nums; this.rand = new Random(); } public int pick(int target) { List<Integer> indices = new ArrayList<>(); for (int i = 0; i < nums.length; i++) { if (nums[i] == target) { indices.add(i); } } int randomIndex = rand.nextInt(indices.size()); return indices.get(randomIndex); }}Complexity
Time
O(N) for each call to `pick(target)`, where N is the number of elements in the array, because we need to scan the entire array. The constructor takes O(1) time.
Space
O(K) for the `pick` method, where K is the number of occurrences of the `target`. This is for the list used to store indices. In the worst-case scenario where all elements are the target, the space complexity becomes O(N). The overall space for the class is O(N) to store the array.
Trade-offs
Pros
Simple to understand and implement.
The constructor is very fast (O(1)) and uses minimal memory beyond storing the array itself.
Cons
The
pickoperation has a time complexity of O(N), which is inefficient if it's called many times.The space complexity for
pickis O(K), where K is the number of occurrences of the target. In the worst case, this can be O(N).
This approach uses Reservoir Sampling to make a random selection in a single pass without using extra space to store all the indices. For each target element encountered, we decide whether to select its index to be our result. The probability of selecting the current index decreases as we find more matches, ensuring that every matched index has an equal final probability of being chosen.
Algorithm
- In the
pick(target)method, initialize a countercountfor the target occurrences to 0 and aresultIndexto -1. - Iterate through the
numsarray with indexi. - If
nums[i]is not equal totarget, continue to the next element. - If
nums[i]is equal totarget, incrementcount. - Generate a random integer between 0 and
count - 1(inclusive). - If this random integer is 0 (which occurs with probability
1/count), updateresultIndexto the current indexi. - After iterating through the whole array, return
resultIndex.
Walkthrough
Reservoir sampling is a clever algorithm that allows for random selection from a list of items of unknown size. Here, we apply it to the indices of the target value. We iterate through the array. When we encounter the k-th occurrence of the target, we replace our current chosen index with the new index with a probability of 1/k. This ensures that after checking all N elements, any of the K occurrences of target is chosen with a probability of 1/K.
For example, for the 1st target, we select its index with probability 1/1. For the 2nd target, we select its index with probability 1/2, keeping the first one with probability 1/2. For the 3rd, we select its index with probability 1/3, keeping the previous choice with probability 2/3. The probability of the first index being the final choice is 1 * (1/2) * (2/3) = 1/3. The probability of the second index being the final choice is (1/2) * (2/3) = 1/3. The probability of the third index being the final choice is 1/3. This generalizes to 1/K for any of the K indices.
import java.util.Random; class Solution { private int[] nums; private Random rand; public Solution(int[] nums) { this.nums = nums; this.rand = new Random(); } public int pick(int target) { int count = 0; int resultIndex = -1; for (int i = 0; i < nums.length; i++) { if (nums[i] == target) { count++; // With probability 1/count, we pick the current index. if (rand.nextInt(count) == 0) { resultIndex = i; } } } return resultIndex; }}Complexity
Time
O(N) for each call to `pick(target)`, as it requires a full scan of the array. The constructor is O(1).
Space
O(1) for the `pick` method, as we only need a couple of variables to store the count and the result. The class itself still requires O(N) space to store the array.
Trade-offs
Pros
The
pickmethod has an excellent space complexity of O(1).It's useful for scenarios with massive datasets or data streams where storing all indices is infeasible.
Cons
The time complexity of
pickis still O(N), which can be slow ifpickis called frequently on a large array.
This approach prioritizes the performance of the pick method by doing pre-computation in the constructor. We can trade space for time by building a hash map that maps each unique number in the input array to a list of all indices where it appears. With this map, a pick operation becomes very efficient: we just look up the target number to get its list of indices and then pick a random one from that list.
Algorithm
- In the constructor, initialize a
HashMapwhere keys are the numbers in the array and values are lists of their indices. - Iterate through the input
numsarray once. - For each element
nums[i], add its indexito the list associated with the keynums[i]in the hash map. - In the
pick(target)method, retrieve the list of indices for the giventargetfrom the hash map. This is an O(1) average time operation. - Generate a random integer between 0 and the size of the list minus 1.
- Return the index from the list at this random position.
Walkthrough
The core idea is to pre-process the data to make subsequent queries faster. The constructor iterates through the nums array and populates a HashMap<Integer, List<Integer>>. The keys of the map are the distinct numbers in nums, and the value for each key is an ArrayList containing all the indices at which that number appears.
Once this map is built, the pick(target) method is simple and fast. It gets the list of indices for the target from the map, which takes O(1) time on average. Then, it generates a random number to select an index from this list. This is the most efficient approach when pick is called multiple times, as the O(N) setup cost is amortized over many fast O(1) queries.
import java.util.ArrayList;import java.util.HashMap;import java.util.List;import java.util.Map;import java.util.Random; class Solution { private Map<Integer, List<Integer>> map; private Random rand; public Solution(int[] nums) { this.map = new HashMap<>(); this.rand = new Random(); for (int i = 0; i < nums.length; i++) { // Get the list for the current number, or create a new one if it doesn't exist map.computeIfAbsent(nums[i], k -> new ArrayList<>()).add(i); } } public int pick(int target) { List<Integer> indices = map.get(target); int randomIndex = rand.nextInt(indices.size()); return indices.get(randomIndex); }}Complexity
Time
O(N) for the constructor to build the map. O(1) on average for each call to `pick(target)`.
Space
O(N) to store the hash map. In the worst case, if all elements are distinct, the map will have N entries, each with a list of one index.
Trade-offs
Pros
The
pickoperation is very fast, with an average time complexity of O(1).This is the most efficient solution when
pickis called many times, as per the problem constraints.
Cons
Requires O(N) extra space to store the hash map.
The constructor has a time complexity of O(N) for pre-computation, which might be a drawback if the constructor needs to be very fast and
pickis called rarely.
Solutions
Solution
class Solution { private int [] nums ; private Random random = new Random (); public Solution ( int [] nums ) { this . nums = nums ; } public int pick ( int target ) { int n = 0 , ans = 0 ; for ( int i = 0 ; i < nums . length ; ++ i ) { if ( nums [ i ] == target ) { ++ n ; int x = 1 + random . nextInt ( n ); if ( x == n ) { ans = i ; } } } return ans ; } } /** * Your Solution object will be instantiated and called as such: * Solution obj = new Solution(nums); * int param_1 = obj.pick(target); */class Solution { public: vector < int > nums ; Solution ( vector < int >& nums ) { this -> nums = nums ; } int pick ( int target ) { int n = 0 , ans = 0 ; for ( int i = 0 ; i < nums . size (); ++ i ) { if ( nums [ i ] == target ) { ++ n ; int x = 1 + rand () % n ; if ( n == x ) ans = i ; } } return ans ; } }; /** * Your Solution object will be instantiated and called as such: * Solution* obj = new Solution(nums); * int param_1 = obj->pick(target); */class Solution : def __init__ ( self , nums : List [ int ]): self . nums = nums 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 # Your Solution object will be instantiated and called as such: # obj = Solution(nums) # param_1 = obj.pick(target)Video walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.