LeetCode 398：随机数索引

在实际开发中，我们经常需要从一组数据中随机选取符合条件的元素，比如在推荐系统里随机挑选一个用户的历史点击视频、在数据库负载均衡中随机选一台可用节点、或者在游戏逻辑中随机生成某种敌人。这道题就是一个小巧但非常有代表性的例子。我们要做的事情是：给定一个整数数组 nums（可能有重复），然后随机返回某个目标值 target 出现的任意一个索引，并且要保证所有满足条件的索引被选中的概率是相等的。

Swift社区

1029人浏览 · 2025-10-14 19:56:07

Swift社区 · 2025-10-14 19:56:07 发布

在这里插入图片描述

文章目录

摘要

在实际开发中，我们经常需要从一组数据中随机选取符合条件的元素，比如在推荐系统里随机挑选一个用户的历史点击视频、在数据库负载均衡中随机选一台可用节点、或者在游戏逻辑中随机生成某种敌人。

这道题就是一个小巧但非常有代表性的例子。
我们要做的事情是：
给定一个整数数组 nums（可能有重复），然后随机返回某个目标值 target 出现的 任意一个索引，并且要保证所有满足条件的索引被选中的 概率是相等的。

听起来简单，其实里面隐藏了很多“概率公平性”的小细节。我们一起来看看怎么优雅地用 Swift 实现它。

描述

题目要求如下：

我们需要实现一个类 Solution，其中包含以下功能：

初始化时传入一个整数数组 nums；
提供方法 pick(target)，从中随机选出一个索引 i，使得 nums[i] == target；
如果 target 有多个位置，要求每个索引被选中的概率相同。

示例

输入：
["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
solution.pick(3) // 随机返回 2、3 或 4 之一

约束条件

1 <= nums.length <= 2 * 10^4
-2^31 <= nums[i] <= 2^31 - 1
题目保证 target 一定存在
pick 最多调用 10^4 次

题解答案

我们有两种典型思路来解这道题：

提前存储所有索引（哈希表法）：
初始化时遍历整个数组，把每个数字出现的索引都存起来。查询时只需要在对应的索引数组中随机取一个就行。
优点是查询速度快，缺点是空间开销较大。
蓄水池抽样法（Reservoir Sampling）：
如果数组特别大，不想额外用哈希表，我们可以一边遍历一边“动态决定”要不要替换结果，这样空间复杂度就是 O(1)。
每次遇到目标值时，以 1/count 的概率更新选中的索引，最终结果保证是等概率的。

我们接下来会分别讲解这两种方法的 Swift 实现，并解释为什么它们都是“公平随机”。

题解代码分析

解法一：哈希表索引法（空间换时间）

import Foundation

class Solution {
    private var indexMap = [Int: [Int]]()
    
    init(_ nums: [Int]) {
        // 预处理：存储每个数字对应的所有索引
        for (i, num) in nums.enumerated() {
            indexMap[num, default: []].append(i)
        }
    }
    
    func pick(_ target: Int) -> Int {
        guard let indices = indexMap[target] else {
            return -1  // 理论上不会出现
        }
        // 随机从数组中选一个
        return indices.randomElement()!
    }
}

解析：

indexMap 用字典来存储：key 是数字，value 是它的所有索引列表；
初始化时一次遍历，构建完整映射；
查询时用 randomElement() 直接随机返回；
由于 Swift 的 randomElement() 是均匀分布的，因此每个索引的概率相等。

这种方式在 nums 规模较小时非常实用，查询非常快，适合多数业务场景。

解法二：蓄水池抽样法（空间优化版）

如果你担心数组太大，不想多占空间，可以用蓄水池抽样来实现 O(1) 空间版本。

import Foundation

class Solution2 {
    private var nums: [Int]
    
    init(_ nums: [Int]) {
        self.nums = nums
    }
    
    func pick(_ target: Int) -> Int {
        var count = 0
        var result = -1
        
        for (i, num) in nums.enumerated() {
            if num == target {
                count += 1
                // 以 1/count 的概率选择当前索引
                if Int.random(in: 0..<count) == 0 {
                    result = i
                }
            }
        }
        
        return result
    }
}

代码讲解：

遍历整个数组；
每次遇到符合 target 的值：
- 递增计数 count；
- 以 1/count 的概率替换当前结果；
最终保证所有符合条件的索引被选中的概率一致。

举个例子：

第一次遇到目标值 → 一定选；
第二次遇到目标值 → 50% 替换；
第三次遇到目标值 → 33% 替换；
这样最终每个索引的选中概率 = 1 / 出现次数，刚好符合要求。

示例测试及结果

let solution = Solution([1, 2, 3, 3, 3])
print(solution.pick(3)) // 随机返回 2、3 或 4
print(solution.pick(1)) // 返回 0
print(solution.pick(3)) // 随机返回 2、3 或 4

let solution2 = Solution2([1, 2, 3, 3, 3])
print(solution2.pick(3)) // 随机返回 2、3 或 4