Random Pick Index

Problem

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 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], [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.

Solution

Method 1 - Linear Search and Reservoir Sampling

We previously discussed this here: [Reservoir Sampling Explained](reservoir-sampling-explained). Now, using the example: {1, 2, 3, 3, 3} with target 3, we aim to select indices 2, 3, and 4, each with a probability of 1/3.

• Index 2: The probability of selection is 1 * (1/2) * (2/3) = 1/3.
• Index 3: The probability of selection is (1/2) * (2/3) = 1/3.
• Index 4: The probability of selection is 1/3.

Code

Java

public class Solution {
    private int[] nums;
    private Random rnd;

    public Solution(int[] nums) {
        this.nums = nums;
        this.rnd = new Random();
    }
    
	public int pick(int target) {
	    int count = 0;
	    for(int i = 0; i < a.length; i++) {
	        if(target == a[i]) {
		        count++;
		    }
	    }
	    int pickIndex = rand.nextInt(count);
	    for(int i = 0; i < a.length; i++) {
	        if(target == a[i]) {
	            if(pickIndex-- == 0) {
		            return i;
		        }
	        }
	    }
	    return 0; // shouldn't come here
	}
}

Complexity

Constructor: Solution(int[] nums)

• ⏰ Time complexity: O(1)
• 🧺 Space complexity: O(1)

Method: int pick(int target)

• ⏰ Time complexity: O(n)
• 🧺 Space complexity: O(1)

Method 2 - Using Map of Number and List of Indices 🏆

We can also do following:

• Create a map where the key is an integer (a number from the input array) and the value is a list of integers (indices where this number appears in the array).
• Generate a random integer i between 0 (inclusive) and the size of currIndices (exclusive).
• Return the element at the randomly selected index from currIndices.

Code

Java

public class Solution {
	private Map<Integer, List<Integer>> indices;
    public Solution(int[] nums) {
	     indices = new HashMap<>();
        for (int i=0; i<nums.length; i++) {
            int num = nums[i];
            indices.putIfAbsent(num, new ArrayList<Integer>());
            indices.get(num).add(i);
        }
    }
    
    public int pick(int target) {
        List<Integer> currIndices = this.indices.get(target);
        int i = (int) (Math.random() * currIndices.size());
        return currIndices.get(i);
    }    
}

Complexity

Constructor: Solution(int[] nums)

• ⏰ Time complexity: O(n)
• 🧺 Space complexity: O(n)

Method: int pick(int target)

• ⏰ Time complexity: O(1)
• 🧺 Space complexity: O(1)