LeetCode.398-随机数索引

题目

给定一个可能含有重复元素的整数数组，要求随机输出给定的数字的索引。 您可以假设给定的数字一定存在于数组中。
注意:
 数组大小可能非常大。 使用太多额外空间的解决方案将不会通过测试。
示例：
int[] nums = new int[] {1,2,3,3,3};
Solution solution = new Solution(nums);
// pick(3) 应该返回索引 2,3 或者 4。每个索引的返回概率应该相等。
solution.pick(3);
// pick(1) 应该返回 0。因为只有nums[0]等于1。
solution.pick(1);

解法一

一道middle难度的题，看到之后思考了一下就有了很直接的思路——先把每个数字的所有下标存起来，然后 pick() 函数得到所有可能下标之后等概率返回其中一个。代码如下：

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

即类初始化的时候遍历数组，用一个 HashMap存每个数字和其所有下标的列表构成的 <key, value> 对，然后对于 pick(int target) 函数，从 hashMap 中获取 target 的下标列表并利用 random 类的 nextInt() 等概率返回其中一个。

解法二

最直接的解答提交通过了，再返回来看题目，注意到了题目里的 “数组大小可能非常大。 使用太多额外空间的解决方案将不会通过测试。”这段话，似乎有点明白了，对于非常大的数组解法一显然是占用了很多额外的空间。想了半天没有更好的思路（还是刷题太少）。然后看了题解才知道这是一个叫做蓄水池抽样的经典问题，对应得解法也很奇妙：
1）遍历数组，计数器 count 记录截至目前遍历到的 target 的个数；
2）对于当前遍历到的 target 的 index，以 \(1/count\) 的概率进行抽样保留。
代码如下：

class Solution {
    int[] nums;
    public Solution(int[] nums) {
        this.nums = new int[nums.length];
        for(int i = 0; i < nums.length; i++){
            this.nums[i] = nums[i];
        }
    }
    
    public int pick(int target) {
        int index = -1;
        Random rand = new Random();
        int count = 0;   // target计数器
        for(int i = 0; i < this.nums.length; i++){
            if(this.nums[i] == target){
                count += 1;
                if(rand.nextInt() % count == 0){  // 以 1/count 的概率对当前 index 进行保留
                    index = i;
                }
            }
        }
        return index;
    }
}

怎么理解呢？
1）首先如果遍历到第一个 target，此时 count=1,对这个索引以1的概率进行保留，如果数组只有这一个 target 那么最后就以概率1返回这个 index1;
2）如果再往后便利有遇到了第二个 target，那么 count+1=2，以 \(1/2\) 的概率保留当前的位置 index2，也就是说有 \(1/2\) 没有替换 index1，这样返回 index1 和 index2 的概率相同；
3）如果再往后便利有遇到了第三个 target，那么 count+1=3，以 \(1/3\) 的概率保留当前的位置 index3，这样 index 的原来的值就有 \(2/3\) 的概率被保留，有由于原来的 index 以相等的概率(\(1/2\))等于 index1 或 index2，这样保留index1 或 index2 的概率就等于了 \(\frac{1}{2}\times \frac{2}{3}=\frac{1}{3}\)。
总结一下就是遍历到第 count 个 target 的时候，以 \(1/count\) 的概率对当前 index 进行保留，\(1-1/count\) 的概率对旧的 index 进行保留，这样归纳可以得到前面每一个 index 的概率都等于 \(1/count\)。

总结

因为刷题少和半天没想出来，看到第二种方法就感觉一个字-秒！
两种方法提交结果来看，解法一耗时竟然比解法二多，有点不太理解，明明 HashMap 和 ArrayList 都可以\(O(1)\)时间获取元素的，而且不需要每次查找索引都便利数组，结果反而慢了。
再说空间大小，看起来好像空间复杂度都是 \(O(n)\)，但其实 HashMap 和 ArrayList 是存的都是实例对象，所以解法二只用了一个额外的 int 数组是节省空间的，可能也是题目的注意核心所在。