题目
给定一个可能含有重复元素的整数数组,要求随机输出给定的数字的索引。 您可以假设给定的数字一定存在于数组中。
注意:
数组大小可能非常大。 使用太多额外空间的解决方案将不会通过测试。
示例:
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} imes 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
数组是节省空间的,可能也是题目的注意核心所在。