LeetCode 398 随机数索引[水塘抽样] HERODING的LeetCode之路

解题思路

第一个方法是哈希+水塘抽样，即用哈希表存储每个num的下标数组，然后pick阶段在每个下标数组中进行水塘抽样，水塘抽样可以将空间复杂度降为O(1)，原理如下:

代码

class Solution {
private:
    unordered_map<int, vector<int>> mp;
public:
    Solution(vector<int>& nums) {
        for(int i = 0; i < nums.size(); i ++) {
            mp[nums[i]].emplace_back(i);
        }
    }
    
    int pick(int target) {
        int n = mp[target].size();
        int index = mp[target][0];
        for(int i = 0; i < n; i ++) {
            if(rand() % (i + 1) == 0) {
                index = mp[target][i];
            }
        }
        return index;
    }
};

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

但是思考一下可以发现，这样做简直就是画蛇添足，因为空间已经被占用了，那么水塘抽样更大可不必，空间复杂度没优化不说还增大了时间复杂度，所以修改一下如下：

class Solution {
private:
    unordered_map<int, vector<int>> mp;
public:
    Solution(vector<int>& nums) {
        for(int i = 0; i < nums.size(); i ++) {
            mp[nums[i]].emplace_back(i);
        }
    }
    
    int pick(int target) {
        int n = mp[target].size();
        int index = mp[target][0];
        for(int i = 0; i < n; i ++) {
            if(rand() % (i + 1) == 0) {
                index = mp[target][i];
            }
        }
        return index;
    }
};

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

那么如何利用水塘抽样的优势呢，可以定义一个vector的引用，直接绑定数组本身，再使用水塘抽样，那么不会耗费额外的空间复杂度了，代码如下：

class Solution {
private:
    vector<int>& nums;
public:
    Solution(vector<int>& nums) : nums(nums){
        //nums = nums;
    } 
    int pick(int target) {
        int n = nums.size();
        int index = 1;
        int ans = 0;
        for(int i = 0; i < n; i ++) {
            if(nums[i] == target) {
                if(rand() % index == 0) {
                    ans = i;
                }
                index ++;
            }
        }
        return ans;
    }
};

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