题目描述

给你一个可能含有 重复元素 的整数数组 nums ，请你随机输出给定的目标数字 target 的索引。你可以假设给定的数字一定存在于数组中。

实现 Solution 类：

Solution(int[] nums) 用数组 nums 初始化对象。
int pick(int target) 从 nums 中选出一个满足 nums[i] == target 的随机索引 i 。如果存在多个有效的索引，则每个索引的返回概率应当相等。

示例：

提示：

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;
    }
}