LeetCode 398. 随机数索引（哈希表/水塘抽样）

---

题目

• 398. 随机数索引

方法一：哈希表

在构造函数中，使用哈希表记录nums中相同元素的下标

在pick函数中，从哈希表中取出target对应的下标列表，然后随机选择其中一个下标并返回

class Solution {
    Map<Integer, List<Integer>> map = new HashMap<>();
    public Solution(int[] nums) {
        for (int i = 0; i < nums.length; i++) {
            if (map.containsKey(nums[i])) {
                map.get(nums[i]).add(i);
            } else {
                List<Integer> list = new ArrayList<>();
                list.add(i);
                map.put(nums[i], list);
            }
        }        
    }
    
    public int pick(int target) {
        Random random = new Random();
        int index = random.nextInt(map.get(target).size());
        return map.get(target).get(index);
    }
}

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

• 时间复杂度：初始化为 $O(n)$，pick 为 $O(1)$，其中 $n$ 是 nums 的长度

• 空间复杂度：$O(n)$。需要 $O(n)$ 的空间存储 $n$ 个下标

方法二：水塘抽样(不定长数据流)

大数据流中的随机抽样问题，即：当内存无法加载全部数据时，如何从包含未知大小的数据流中随机选取k个数据，并且要保证每个数据被抽取到的概率相等

当 $k=1$时：

假设数据流含有$N$个数，如果要保证所有的数被抽到的概率相等，那么每个数抽到的概率应该为 $\frac{1}{N}$.

可以这样做：

• 遇到第$1$个数 $n_1$ 的时候，保留它， $p(n_1)=1$
• 遇到第$2$个数 $n_2$ 的时候，以 $\frac{1}{2}$ 的概率保留它，那么 $p(n_1)=1\times \frac{1}{2}=\frac{1}{2}$，$p(n_2)=\frac{1}{2}$
• 遇到第$3$个数 $n_3$ 的时候，以 $\frac{1}{3}$ 的概率保留它，那么 $p(n_1)=p(n_2)=\frac{1}{2}\times (1-\frac{1}{3})=\frac{1}{3}$，$p(n_3)=\frac{1}{3}$
• $\cdots$
• 遇到第$i$个数 $n_i$ 的时候，以 $\frac{1}{i}$ 的概率保留它，那么 $p(n_1)=p(n_2)=p(n_3)=\cdots =p(n_{i-1})=\frac{1}{i-1}\times (1-\frac{1}{i})=\frac{1}{i}$，$p(n_i)=\frac{1}{i}$

可以看出，对于$k=1$的情况，可以制定这样简单的抽样策略：

数据流中第$i$个数被保留的概率为 $\frac{1}{i}$ 。只要采取这种策略，只需要遍历一遍数据流就可以得到采样值，并且保证所有数被选取的概率均为 $\frac{1}{N}$。

当$k>1$时：

仍然假设数据流中含有$N$个数，那么要保证所有的数被抽到的概率相等，每个数被选取的概率必然为 $\frac{k}{N}$。

• 对于前$k$个数 $n_1,n_2,\cdots ,n_k$，保留下来，则 $p(n_1)=p(n_2)=\cdots =p(n_k)=1$（下面连等采用 $p(n_{1:k})$ 的形式）
• 对于第$k+1
  $个数 $n_{k+1}$，以 $\frac{k}{k+1}$ 的概率保留它（这里只是指本次被保留下来），那么前$k$个数中的$n_r(r\in 1:k)$被保留的概率可以这样表示：$p(n_r被保留)=p(上一轮n_r被保留)\times (p(n_{k+1}被丢弃)+p(n_{k+1}被保留)\times p(n_r未被替换))$，即$p(n_{1:k})=\frac{1}{k+1}+\frac{k}{k+1}\times \frac{k-1}{k}=\frac{k}{k+1}$
• 对于第$k+2$个数 $n_{k+2}$，以 $\frac{k}{k+2}$ 的概率保留它（这里只是指本次被保留下来），那么前$k$个被保留下来的数中的$n_r(r\in 1:k)$被保留的概率为$p(n_{1:k})=\frac{k}{k+1}\times (\frac{2}{k+2}+\frac{k}{k+2}\times \frac{k-1}{k})=\frac{k}{k+2}$
• $\cdots$
• 对于第$i（i>k）$个数 $n_i$ ，以 $\frac{k}{i}$ 的概率保留它，前$i-1$个数中的 $n_r(r\in 1:i-1)$ 被保留的概率为 $p(n_{1:k})=\frac{k}{i-1}\times (\frac{i-k}{i}+\frac{k}{i}\times \frac{k-1}{k})=\frac{k}{i}$

这样，我们可以制订策略：

对于前$k$个数，全部保留，对于第$i（i>k）$个数，以 $\frac{k}{i}$ 的概率保留第$i$个数，并以 $\frac{1}{k}$ 的概率与前面已选择的$k$个数中的任意一个替换。

class Solution {
    Random random = new Random();
    int[] nums;
    public Solution(int[] _nums) {
        nums = _nums;
    }
    public int pick(int target) {
        int n = nums.length, ans = 0;
        for (int i = 0, cnt = 0; i < n; i++) {
            if (nums[i] == target) {
                cnt++;
                if (random.nextInt(cnt) == 0) ans = i;
            }
        }
        return ans;
    }
}

• 时间复杂度：初始化的复杂度为 $O(1)$；pick 操作的复杂度为 $O(n)$
• 空间复杂度：$O(n)$

Reference

• 水塘抽样（Reservoir Sampling）

• 一题双解 :「哈希表预处理」&「蓄水池抽样」