Random Pick Index

https://leetcode.com/problems/random-pick-index/description/

Given an array of integers 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.

Note:

The array size can be very large. Solution that uses too much extra space will not pass the judge.

Example:

int[] nums = new int[] {1,2,3,3,3};

Solution solution = new Solution(nums);

// pick(3) should return either index 2, 3, or 4 randomly. Each index should have equal probability of returning.

solution.pick(3);

// pick(1) should return 0. Since in the array only nums[0] is equal to 1.

solution.pick(1);

Thoughts

给一个数字组成的stream，让随机返回给定一个数的index。给定数有很多重复，不能把它的所有可能的index都装入内存。水塘抽样适用于当内存无法加载全部数据时，如何从包含未知大小的数据流中随机选取k个数据，并且要保证每个数据被抽取到的概率相等。 之前一般是空间换时间, 现在是时间换空间。 从未知个数的nums中抽k个sample, 遍历nums, 前k个全部保留, 后面的以k/i保留, 并以1/k概率替换之前选出来的k个。 当k=1时, 保留第一个, 之后每个以1/i概率替换。 PENG：水塘抽样（Reservoir Sampling）​zhuanlan.zhihu.com

{1,2,3,3,3} with target 3, you want to select 2,3,4 with a probability of 1/3 each. 2 : It's probability of selection is 1 (1/2) (2/3) = 1/3 3 : It's probability of selection is (1/2) * (2/3) = 1/3 4 : It's probability of selection is just 1/3

Code

/*
 * @lc app=leetcode id=398 lang=cpp
 *
 * [398] Random Pick Index
 *
 * https://leetcode.com/problems/random-pick-index/description/
 *
 * algorithms
 * Medium (50.92%)
 * Likes:    293
 * Dislikes: 495
 * Total Accepted:    61.5K
 * Total Submissions: 119.8K
 * Testcase Example:  '["Solution","pick"]\n[[[1,2,3,3,3]],[3]]'
 *
 * Given an array of integers 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.
 * 
 * Note:
 * The array size can be very large. Solution that uses too much extra space
 * will not pass the judge.
 * 
 * Example:
 * 
 * 
 * int[] nums = new int[] {1,2,3,3,3};
 * Solution solution = new Solution(nums);
 * 
 * // pick(3) should return either index 2, 3, or 4 randomly. Each index should
 * have equal probability of returning.
 * solution.pick(3);
 * 
 * // pick(1) should return 0. Since in the array only nums[0] is equal to 1.
 * solution.pick(1);
 * 
 * 
 */
class Solution {
public:
    vector<int> _nums;
    Solution(vector<int>& nums) {
        _nums = nums;
    }
    
    int pick(int target) {
        int k = 1, count = 0;
        vector<int> res(k);
        for (int i = 0; i < _nums.size(); ++i) {
            if (_nums[i] != target) continue;
            if (count < k) {
                res[count++] = i;
            } else {
                const int r = rand() % (++count);
                if (r < k) res[r] = i;
            }
        }
        return res[k - 1];      
    }
};

/**
 * Your Solution object will be instantiated and called as such:
 * Solution* obj = new Solution(nums);
 * int param_1 = obj->pick(target);
 */

class Solution {
    int[] nums;
    Random rand;
    public Solution(int[] nums) {
        this.nums = nums;
        rand = new Random();
    }

    public int pick(int target) {
        int count = 0;
        int k = 1;
        int[] res = new int[k];

        for (int i = 0; i < nums.length; i++) {
            if (nums[i] == target) {
                if (count < k) {
                    res[count++] = i;
                } else {
                    int r = rand.nextInt(++count);
                    if (r < k) {
                        res[r] = i;
                    }
                }
            }
        }
        return res[0];
    }
}

/**
 * Your Solution object will be instantiated and called as such:
 * Solution obj = new Solution(nums);
 * int param_1 = obj.pick(target);
 */

Analysis

每次pick的时间复杂度O(N).