Clean Java solution with explanation of resevior sampling


  • 0
    J
    public int pick(int target) {
        int candidate = -1;
        int count = 0;
        int i = 0;
        while (i < nums.length) {
            if (nums[i] == target) {
                if (candidate == -1) {
                    candidate = i;
                    count++;
                } else {
                    count++;
                    int random = (int) (Math.random() * count);
                    if (random == 0) {
                        candidate = i;
                    }
                }
            }
            
            i++;
        }
        return candidate;
    }
    

    the easy solution is to use hashMap and list to store the index of each element. And then random select an element in the list. However, got TLE or MLE sometimes.

    The accepted solution is to use resevior sampling. First, we need an pool of eligible elements (size == 1 in my solution). And count indicates the number of targets so far when we get to nums[i]. Second, we get a random number and if it equals to 0, we set the candidate to i. The possibility is 1/count.

    More generally, we can have a pool size of k. So when we can get to element i, the probability that element i can be placed in the pool is k/i. And from the i + 1 to n - 1, the probability that element i will not be replaced on is (1 - 1 / (i + 1)) * (1 - 1 / (i + 2)) * ..... * (1 - 1 / n) so the final result is k/n.


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.