[leetcode] 398. 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);

分析

题目的意思是：等概率的返回数组里面的值。

• 水塘抽样算法。水塘抽样是一系列的随机算法，其目的在于从包含n个项目的集合S中选取k个样本，其中n为一很大或未知的数量，尤其适用于不能把所有n个项目都存放到内存的情况。

从S中抽取k项放入「水塘」中
对于每个S[j]项（j ≥ k）：
   随机产生一个范围从0到j的整数r
   若 r < k 则把水塘中的第r项换成S[j]项

• 定义两个变量，计数器cnt和返回结果res，遍历整个数组，如果数组的值不等于target，直接跳过；如果等于target，计数器加1，然后我们在[0,cnt)范围内随机生成一个数字，如果这个数字是0，我们将res赋值为i即可

代码

class Solution {
private:
    vector<int> v;
public:
    Solution(vector<int> nums):v(nums) {
        
    }
    
    int pick(int target) {
        int cnt=0,res=-1;
        for(int i=0;i<v.size();i++){
            if(v[i]!=target) continue;
            ++cnt;
            if(rand()%cnt==0) res=i;
        }
        return res;
    }
};

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

参考文献

[LeetCode] Random Pick Index 随机拾取序列
水塘抽样