LC 398. 随机数索引

题目描述

这是 LeetCode 上的 398. 随机数索引 ，难度为 中等。

给定一个可能含有重复元素的整数数组，要求随机输出给定的数字的索引。

您可以假设给定的数字一定存在于数组中。

注意：数组大小可能非常大，使用太多额外空间的解决方案将不会通过测试。

示例:

1
2
3
4
5
6
7
8
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);

提示：

  • $1 <= nums.length <= 2 \times 10^4$
  • $-2^{31} <= nums[i] <= 2^{31} - 1$
  • target 确保存在于 nums 中
  • 最多调用 $10^4$ 次的 pick

哈希表 + 预处理（定长数据流）

为了方便，我们令 nums 的长度为 $n$，利用 $n$ 的数据范围为 $2 \times 10^4$，且完整的数组为初始化时已给出，我们可以通过使用「哈希表 + 预处理」的方式进行求解。

具体的，在构造函数传入 nums 时，遍历 nums 并存储每个 $nums[i]$ 对应的下标集合，即使用哈希表以 $nums[i]$ 为键，下标集合 List 作为值进行存储。

在 pick 操作时，通过 $O(1)$ 的复杂度取出所有 $nums[i] = target$ 的集合下标，再随机一个下标进行返回。

Java 代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
class Solution {
    Random random = new Random();
    Map<Integer, List<Integer>> map = new HashMap<>();
    public Solution(int[] nums) {
        int n = nums.length;
        for (int i = 0; i < n; i++) {
            List<Integer> list = map.getOrDefault(nums[i], new ArrayList<>());
            list.add(i);
            map.put(nums[i], list);
        }
    }
    public int pick(int target) {
        List<Integer> list = map.get(target);
        return list.get(random.nextInt(list.size()));
    }
}

C++ 代码：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
class Solution {
public:
    unordered_map<int, vector<int>> map;
    Solution(vector<int>& nums) {
     int n = nums.size();
        for (int i = 0; i < n; i++) {
            map[nums[i]].push_back(i);
        }
    }
    int pick(int target) {
        vector<int>& list = map[target];
        return list[rand() % list.size()];
    }
};

Python 代码：
1
2
3
4
5
6
7
8
9
10
11
class Solution:
    def __init__(self, nums: List[int]):
        self.mapping = {}
        n = len(nums)
        for i in range(n):
            if nums[i] not in self.mapping:
                self.mapping[nums[i]] = []
            self.mapping[nums[i]].append(i)

    def pick(self, target: int) -> int:
        return random.choice(self.mapping.get(target, []))

  • 时间复杂度：初始化的复杂度为 $O(n)$；pick 操作的复杂度为 $O(1)$
  • 空间复杂度：$O(n)$

蓄水池抽样（不定长数据流）- TLE

虽然对于本题而言，蓄水池抽样方法（pick 操作复杂度 $O(n)$ 的）会面临 TLE。

但当 nums 并不是在初始化时完全给出，而是持续以「流」的形式给出，且数据流的很长，不便进行预处理的话，我们只能使用「蓄水池抽样」的方式求解。

不了解「蓄水池抽样」的同学可以看前置 🧀 : 多语言入门「蓄水池抽样」知识点

具体的，我们在每次 pick 时对流进行遍历，由于数据流很大，我们不能在遍历过程中使用诸如数组的容器存储所有满足条件的下标，只能对于每个 $nums[i] = target$ 执行「是否要将 $i$ 作为最新答案候选」的操作。

假设共有 $m$ 个下标满足 $nums[i] = target$，我们需要做到以 $\frac{1}{m}$ 概率返回任一坐标。

我们规定当遇到第 $k$ 个满足 $nums[i] = target$ 的下标时，执行一次 $[0, k)$ 的随机操作，当随机结果为 $0$ 时（发生概率为 $\frac{1}{k}$），我们将该坐标作为最新的答案候选。

当对每一个 $nums[i] = target$ 的下标都进行上述操作后，容易证明每一位下标返回的概率均为 $\frac{1}{m}$。

假设最后返回的是第 $k$ 个满足条件的下标，发生概率为 = 第 $k$ 个下标被候选的概率 $\times$ 后面 $k + 1$ 到 $m$ 个下标不被候选的概率 = $\frac{1}{k} \times (1 - \frac{1}{k + 1}) \times … \times (1 - \frac{1}{m})$ = $\frac{1}{m}$ 。

Java 代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
class Solution {
    Random random = new Random();
    int[] nums;
    public Solution(int[] _nums) {
        nums = _nums;
    }
    public int pick(int target) {
        int n = nums.length, ans = 0;
        for (int i = 0, cnt = 0; i < n; i++) {
            if (nums[i] == target) {
                cnt++;
                if (random.nextInt(cnt) == 0) ans = i;
            }
        }
        return ans;
    }
}

C++ 代码：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution {
public:
    vector<int> nums;
    mt19937 random;
    Solution(vector<int>& _nums) : nums(_nums) {
        srand(time(nullptr));
        random.seed(rand());
    }
    int pick(int target) {
        int ans = -1, cnt = 0;
        for (size_t i = 0; i < nums.size(); i++) {
            if (nums[i] == target) {
                cnt++;
                if (uniform_int_distribution<>(0, cnt - 1)(random) == 0) {
                    ans = i;
                }
            }
        }
        return ans;
    }
};

Python 代码：
1
2
3
4
5
6
7
8
9
10
11
class Solution:
    def __init__(self, nums: List[int]):
        self.nums = nums
    def pick(self, target: int) -> int:
        ans, cnt = 0, 0
        for i in range(len(self.nums)):
            if self.nums[i] == target:
                cnt += 1
                if random.randint(0, cnt - 1) == 0:
                    ans = i
        return ans

  • 时间复杂度：初始化的复杂度为 $O(1)$；pick 操作的复杂度为 $O(n)$
  • 空间复杂度：$O(n)$

最后

这是我们「刷穿 LeetCode」系列文章的第 No.398 篇，系列开始于 2021/01/01，截止于起始日 LeetCode 上共有 1916 道题目，部分是有锁题，我们将先把所有不带锁的题目刷完。

在这个系列文章里面，除了讲解解题思路以外，还会尽可能给出最为简洁的代码。如果涉及通解还会相应的代码模板。

为了方便各位同学能够电脑上进行调试和提交代码，我建立了相关的仓库：https://github.com/SharingSource/LogicStack-LeetCode 。

在仓库地址里，你可以看到系列文章的题解链接、系列文章的相应代码、LeetCode 原题链接和其他优选题解。