Java&C++题解与拓展——leetcode398.随机数索引【水塘抽样学习】

每日一题做题记录，参考官方和三叶的题解

文章目录

• 题目要求
• 思路一：模拟、哈希表
• • Java
  • C++
• 思路二：水塘抽样（蓄水池抽样）
• • Java
  • C++
• 总结

题目要求

思路一：模拟、哈希表

把数组内容整理一下放哈希表，然后从哈希表取值随机返回。
哈希表存的内容是数组元素和它对应的所有下标。

Java

class Solution {
    Map<Integer, List<Integer>> map;
    Random ran;

    public Solution(int[] nums) {
        map = new HashMap<Integer, List<Integer>>();
        ran = new Random();
        for(int i = 0; i < nums.length; ++i) {
            map.putIfAbsent(nums[i], new ArrayList<Integer>()); // 相同元素放一起
            map.get(nums[i]).add(i); // 存下标
        }
    }
    
    public int pick(int target) {
        List<Integer>idx = map.get(target); // 取下标
        return idx.get(ran.nextInt(idx.size()));
    }
}

• 时间复杂度：初始化为 O ( n ) O(n) O(n)，pick函数为 O ( 1 ) O(1) O(1)
• 空间复杂度： O ( n ) O(n) O(n)

C++

class Solution {
    unordered_map<int, vector<int>> map;
public:
    Solution(vector<int>& nums) {
        for(int i = 0; i < nums.size(); ++i)
            map[nums[i]].push_back(i); // 相同元素下标放一起
    }
    
    int pick(int target) {
        auto &idx = map[target]; // 取下标
        return idx[rand() % idx.size()];
    }
};

• 时间复杂度：初始化为 O ( n ) O(n) O(n)，pick函数为 O ( 1 ) O(1) O(1)
• 空间复杂度： O ( n ) O(n) O(n)

思路二：水塘抽样（蓄水池抽样）

降低空间复杂度，边读边取，适用于长长长文件读取处理。

• 遍历 n u m s nums nums，每次遇到 t a r g e t target target元素都选择性更新结果。
• 设当前为第 c n t cnt cnt次，选择方法为产生 [ 0 , c n t ) [0,cnt) [0,cnt)内的一个随机整数 r a n ran ran：
  • 若 r a n = 0 ran=0 ran=0，更新结果为当前元素在数组中的下标（不是 c n t cnt cnt）；
  • 若 r a n ≠ 0 ran\ne 0 ran​=0：不更新结果。

这个选择方法是怎么保证返回每个下标概率相同的呢？
P ( 第 i 个 t a r g e t 元 素 下 标 成 为 结 果 ) \quad P(第i个target元素下标成为结果) P(第i个target元素下标成为结果)
P ( r a n i = 0 ) × P ( r a n i + 1 ≠ 0 ) × ⋯ × P ( r a n k ≠ 0 ) \quad P(ran_i=0)\times P(ran_{i+1}\ne 0)\times \dots\times P(ran_k\ne 0) P(rani​=0)×P(rani+1​​=0)×⋯×P(rank​​=0)
= 1 i × ( 1 − 1 i + 1 ) × ⋯ × ( 1 − 1 k ) =\frac{1}{i}\times(1-\frac{1}{i+1})\times\dots\times(1-\frac{1}{k}) =i1​×(1−i+11​)×⋯×(1−k1​)
= 1 i × i i + 1 × ⋯ × k − 1 k =\frac{1}{i}\times\frac{i}{i+1}\times\dots\times\frac{k-1}{k} =i1​×i+1i​×⋯×kk−1​
= 1 k =\frac{1}{k} =k1​
注： r a n i ran_i rani​指第 i i i轮中选择的随机数。

Java

class Solution {
    int[] nums;
    Random ran;

    public Solution(int[] nums) {
        this.nums = nums;
        ran = new Random();
    }
    
    public int pick(int target) {
        int res = 0;
        for(int i = 0, cnt = 0; i < nums.length; ++i) {
            if(nums[i] == target) {
                ++cnt; // 第cnt个target
                if(ran.nextInt(cnt) == 0)
                    res = i;
            }
        }
        return res;
    }
}

• 时间复杂度：初始化为 O ( 1 ) O(1) O(1)，pick函数为 O ( n ) O(n) O(n)
• 空间复杂度： O ( 1 ) O(1) O(1)

C++

class Solution {
    vector<int> &nums;
public:
    Solution(vector<int>& nums) : nums(nums) {}
    
    int pick(int target) {
        int res;
        for(int i = 0, cnt = 0; i < nums.size(); ++i) {
            if(nums[i] == target) {
                ++cnt; // 第cnt个target
                if(rand() % cnt == 0)
                    res = i;
            }
        }
        return res;
    }
};

• 时间复杂度：初始化为 O ( 1 ) O(1) O(1)，pick函数为 O ( n ) O(n) O(n)
• 空间复杂度： O ( 1 ) O(1) O(1)

总结

快乐题目+1，学了新的抽样方法，可以用来处理不定长的巨大数据流，还能保证对每个数的抽取概率一致。

欢迎指正与讨论！