398. 随机数索引
题目描述
给你一个可能含有 重复元素 的整数数组 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 次
答案
我的代码
class Solution { int[] arr ; public Solution(int[] nums) { arr = nums; } public int pick(int target) { List<Integer> list = new ArrayList<>(); for (int i = 0; i < arr.length; i++) { if (arr[i]==target){ list.add(i); } } // 生成随机数 Random random = new Random(); // 返回结果 return list.get(random.nextInt(list.size())); } }
官网答案
方法一:哈希表
如果不考虑数组的大小,我们可以在构造函数中,用一个哈希表 pos 记录 nums 中相同元素的下标。
对于 pick 操作,我们可以从 pos 中取出 target 对应的下标列表,然后随机选择其中一个下标并返回。
class Solution { Map<Integer, List<Integer>> pos; Random random; public Solution(int[] nums) { pos = new HashMap<Integer, List<Integer>>(); random = new Random(); for (int i = 0; i < nums.length; ++i) { pos.putIfAbsent(nums[i], new ArrayList<Integer>()); pos.get(nums[i]).add(i); } } public int pick(int target) { List<Integer> indices = pos.get(target); return indices.get(random.nextInt(indices.size())); } }
复杂度分析
时间复杂度:初始化为 O(n),pick 为 O(1),其中 n 是 nums 的长度。
空间复杂度:O(n)。我们需要 O(n) 的空间存储 n 个下标。
方法二:水塘抽样
class Solution { int[] nums; Random random; public Solution(int[] nums) { this.nums = nums; random = new Random(); } public int pick(int target) { int ans = 0; for (int i = 0, cnt = 0; i < nums.length; ++i) { if (nums[i] == target) { ++cnt; // 第 cnt 次遇到 target if (random.nextInt(cnt) == 0) { ans = i; } } } return ans; } }
时间复杂度:初始化为 O(1),pick 为 O(n),其中 n 是 nums 的长度。
空间复杂度:O(1)。我们只需要常数的空间保存若干变量。