[Leetcode 398] Random Pick Index

作者 Shilei Tian 日期 2017-01-17
[Leetcode 398] Random Pick Index

题目要求

Given an array of integers with possible duplicates, randomly output the index of a given target number. You can assume that the given target number must exist in the array.

Note:

The array size can be very large. Solution that uses too much extra space will not pass the judge.

Example:

1
2
3
4
5
6
7
8
9
10
11
12
> int[] nums = new int[] {1,2,3,3,3};
>
> Solution solution = new Solution(nums);
>
> // pick(3) should return either index 2, 3, or 4 randomly. Each index should have equal probability of returning.
>
> solution.pick(3);
>
> // pick(1) should return 0. Since in the array only nums[0] is equal to 1.
>
> solution.pick(1);
>

解题思路

看到这道题后,能立马想到一个很简单的思路,把目标出现的下标都存下来,然后在随机地在这个数组里面选一个就行了,但是被题目要求给吓住了,它说给定的数组可能会非常大,所以我在想是不是这种方式不会通过。

那就换个思路,先将数组排序,然后找到目标出现的范围,再在这个范围内随机选取一个位置,但是又感觉 $O(n\log n)$ 的算法可能会超时。

算了,先看看本题的 Tag 是什么吧,诶,一点,是 Reservoir Sampling,这从来没听说过啊…看了下 Wiki,才明白是怎么回事儿,读到第二部分 Algorithm R 就可以仿造它来做这道题了。

下面我们来说这道题的思路,初始序列没有顺序,我们可以把它当成 Wiki 里面所谓的一个流,设置一个计数变量 count,从第一个元素开始起进行扫描,检查当前元素是否和目标元素相等,如果相等,就对 count 进行自增,与此同时,我们随机生成一个数 $num \in \lbrack 0,count)$ 代表我们会选择的数字,检查 num,如果 num = count - 1,就说明我们要选的数就是最后新加入的这个;如果不是的话,就说明我们要保留之前选择的。

为什么这样做可以得到正确的结果,证明过程我们就不在这里赘述,Wiki 上讲述得很清楚,这个其实就和抓阄没有先后顺序是类似的,概率都是等同的。

具体的代码如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution {
public:
Solution(vector<int> n) {
nums = n;
}

int pick(int target) {
int count = 0, res = -1;
for (int i = 0; i < nums.size(); ++i) {
if (nums[i] != target) {
continue;
}
if (++count == 1) {
res = i;
} else if (rand() % count == count - 1) {
res = i;
}
}
return res;
}

private:
vector<int> nums;
};

这样下来,时间复杂度就是 $O(n)$,还是很快的。

Your runtime beats 99.75% of cpp submissions.