Solution
The question can be easily solved using a map which stores each number and its positions. Once we build the map, picking a random position for the target number is an O(1) operation as we only have to pick a random number from map[target]. But this may require huge amount of memory if the nums array is very large (as mentioned in the question). Although this solution is optimal in terms of time [O(n) preprocessing, O(1) picking], it requires O(n) space.
We can trade off space for time as asked by the question. We will employ a standard technique, reservoir sampling, to solve this question. Please read about reservoir sampling before proceeding.
We only need a reservoir of size 1 which will store the index we want to return. The reservoir may have any index of the target number in the nums array. Reservoir sampling guarantees that each index has equal probability to be selected.
class Solution {
// reservoir sampling
// https://www.geeksforgeeks.org/reservoir-sampling/
public:
random_device rd;
vector<int> nums;
Solution(vector<int>& nums) {
this->nums = nums;
}
int pick(int target) {
int res;
int count = 0;
for (int i=0; i<nums.size(); i++) {
if (nums[i] != target)
continue;
count++;
int rand = rd()%(count);
if (rand==0) {
res = i;
}
}
return res;
}
};
Consider an example:
nums=[1,5,5,6,5,7,9,0] and the target is 5. 5 is at indices 1,2, and 4.
At index 1:
count = 1 so rand can only have the value 0
So res = 1 (probability = 1)
At index 2:
count = 2 so rand can have value {0, 1}
We only change res when rand = 0 so probability of keeping res as 1 is (1/2)
At index 4:
count = 3 so rand can have value {0, 1, 2}
We only change res when rand = 0 so probability of keeping res as 1 is (2/3)
So, having res=1 finally has probability 1*1/2*2/3 = 1/3. You can do the same analysis for rest of the indices.
Time complexity: O(n) for picking Space complexity: O(1)