持续创作，加速成长！这是我参与「掘金日新计划 · 6 月更文挑战」的第2天，点击查看活动详情
题目描述
给你一个可能含有 重复元素 的整数数组 nums ，请你随机输出给定的目标数字 target 的索引。你可以假设给定的数字一定存在于数组中。
实现 Solution 类：
Solution(int[] nums) 用数组 nums 初始化对象。
int pick(int target) 从 nums 中选出一个满足 nums[i] == target 的随机索引 i 。如果存在多个有效的索引，则每个索引的返回概率应当相等。
示例：
输入
["Solution", "pick", "pick", "pick"]
[[[1, 2, 3, 3, 3]], [3], [1], [3]]
输出
[null, 4, 0, 2]
解释
- Solution solution = new Solution([1, 2, 3, 3, 3]);
- solution.pick(3); // 随机返回索引 2, 3 或者 4 之一。每个索引的返回概率应该相等。
- solution.pick(1); // 返回 0 。因为只有 nums[0] 等于 1 。
- solution.pick(3); // 随机返回索引 2, 3 或者 4 之一。每个索引的返回概率应该相等。
提示：
- 1 <= nums.length <= 2 * 104
- -231 <= nums[i] <= 231 - 1
- target 是 nums 中的一个整数
- 最多调用 pick 函数 104 次
思路
本题其实是“水塘抽样算法”的一个变形，所以先介绍一下这个有点意思的“水塘抽样算法”。
问题是这样：现在有一个长度未知的链表，需要设计一种算法，每次随机返回链表中的一个元素，使得每个元素返回的概率是相等的，且最多只能遍历这个链表1次。
常规的随机思路是数组，我们先求出数组的长度len，然后使用random.nextInt(len)获取范围0~len-1的下标，此时每个下标的概率都是1/len，然后返回num[r]即可。但是我们套用到上面的问题中就发现，如果要知道链表的长度len，先要遍历一遍链表，假设随机到的下标是r，那么又要从头结点开始向后走r次，才能得到要返回的节点，但是这样就不符合只能遍历1次这个限制了。
水塘抽样算法就是解决这个问题的一种设计：使用一个变量保存返回的元素，遍历到第i个元素（下标为i-1）时，以1/i的概率将这个元素保存到待返回的变量中，那么整体上来说，每个元素被返回的概率都是1/len。
我们来证明一下算法的正确性，如果最后返回的是第i个元素
- 跟之前的i-1个元素是否被选中放到待返回的变量中无关
- 在第i个元素做选择时，必然被选中，概率为
1/i - 在第i个元素后做选择时，都不被选中，概率是为(1-1/i+1)(1-1/i+2)...*(1-1/len) 综合下来，第i个元素被选中的概率为
水塘抽样算法的基础代码
int getRandom(ListNode head) {
Random r = new Random();
int i = 0, res = 0;
ListNode p = head;
while (p != null) {
if (r.nextInt(++i) == 0) {
res = p.val;
}
p = p.next;
}
return res;
}
回到本题，本题要随机返回等于target的数组下标，那么我们只要增加一步处理，就是先选择出等于target的元素，然后使用“水塘抽样算法”，边判断等于边随机抽取即可。
Java版本代码
class Solution {
private int[] nums;
private Random random;
private int len;
public Solution(int[] nums) {
this.nums = nums;
len = nums.length;
random = new Random();
}
public int pick(int target) {
int ans = -1;
for (int i = 0, cnt = 0; i < len; i++) {
if (target == nums[i]) {
cnt++;
if (random.nextInt(cnt) == 0) {
ans = i;
}
}
}
return ans;
}
}
/**
* Your Solution object will be instantiated and called as such:
* Solution obj = new Solution(nums);
* int param_1 = obj.pick(target);
*/
