Problem explanation

Given an array of integers with possible duplicate elements, we are to generate the index of such element uniformly at random. Let's consider an array nums = [1,2,3,3,3], here, if we call the function pick(1) we should return the index 0 (Because only index 0 is equal to 1) and for pick(3) it should return index 2, 3 or 4.

Our goal is to create an algorithm that would return the index of the elements at random without using too much space. Therefore, we can't store all indices and then generate one uniformly at random. We will have to do it on the fly.

Approach

This problem can be solved using the Reservoir Sampling technique. The idea of Reservoir Sampling is simple whenever we receive a sample we replace the previous sample with a probability of 1 / cnt where cnt is the count of the sample we have seen till now.

Examples

Consider an input nums = [1, 2, 3, 3, 3]. Creating a class Solution as below

Solution solution = new Solution(nums);

Now, if we call pick(3)

The function will process elements of the array one by one and if the element is equal to target, it will keep track of the count (cnt) of occurrence of the target elements and replace the index with the current index with a probability of 1 / cnt.

Python solution

1
2python
3import random
4class Solution:
5
6    def __init__(self, nums):
7        self.nums = nums
8
9    def pick(self, target):
10        res = -1
11        count = 0
12        for i, x in enumerate(self.nums):
13            if x == target:
14                if random.randint(0, count) == 0:
15                    res = i
16                count += 1
17        return res

Java solution

1
2java
3class Solution {
4    private int[] nums;
5
6    public Solution(int[] nums) {
7       this.nums = nums;
8    }
9
10    public int pick(int target) {
11        int result = -1;
12        int count = 0;
13        for (int i = 0; i < nums.length; i++) {
14            if (nums[i] != target)
15                continue;
16            if (rand.nextInt(++count) == 0)
17                result = i;
18        }
19
20        return result;
21    }
22}

JavaScript solution

1
2javascript
3var Solution = function(nums) {
4    this.nums = nums;
5};
6
7Solution.prototype.pick = function(target) {
8    var result = -1;
9    var count = 0;
10    for (var i = 0; i < this.nums.length; i++) {
11        if (this.nums[i] !== target)
12            continue;
13        if ((Math.floor(Math.random() * (++count))) === 0)
14            result = i;
15    }
16    return result;
17};

C++ solution

1
2cpp
3class Solution {
4public:
5    vector<int> nums;
6
7    Solution(vector<int>& nums) {
8        this->nums = nums;
9    }
10
11    int pick(int target) {
12        int result = -1;
13        int count = 0;  
14        for(int i = 0; i < nums.size(); i++){
15            if(nums[i] != target)
16                continue;
17            if(rand() % (++count) == 0)
18                result = i;
19        }
20        return result;
21    }
22};

C# solution

1
2csharp
3public class Solution {
4    private int[] nums;
5    private Random rand;
6
7    public Solution(int[] nums) {
8        this.nums = nums;
9        this.rand = new Random();
10    }
11
12    public int Pick(int target) {
13        int result = -1;
14        int count = 0;
15        for (int i = 0; i < nums.Length; i++) {
16            if (nums[i] != target) 
17                continue;
18            if (rand.Next(++count) == 0) 
19                result = i;
20        }
21        return result;
22    }
23}

Golang solution

1
2go
3import "math/rand"
4
5type Solution struct {
6    nums []int
7}
8
9func Constructor(nums []int) Solution {
10    return Solution{nums: nums}
11}
12
13func (this *Solution) Pick(target int) int {
14    count := 0
15    result := -1
16    for i, n := range this.nums {
17        if n == target {
18            count++
19            if rand.Intn(count) == 0 {
20                result = i
21            }
22        }
23    }
24    return result
25}

In Go, the math/rand library is used to generate our random number.

Conclusion

The Reservoir Sampling technique is a very useful way of solving problems like these that require randomness but with limited space usage. Such a technique keeps the space complexity at O(1) and time complexity at O(n) where n is the number of integers. It allows us to efficiently deal with large datasets without requiring large amounts of storage. You are now equipped to handle such problems by simply counting the occurrences and using a random integer generator to select an index uniformly at random.