LeetCode 398:随机数索引
文章目录
- 摘要
- 描述
- 示例
- 约束条件
- 题解答案
- 题解代码分析
- 解法一:哈希表索引法(空间换时间)
- 解法二:蓄水池抽样法(空间优化版)
- 示例测试及结果
- 时间复杂度
- 空间复杂度
- 总结
摘要
在实际开发中,我们经常需要从一组数据中随机选取符合条件的元素,比如在推荐系统里随机挑选一个用户的历史点击视频、在数据库负载均衡中随机选一台可用节点、或者在游戏逻辑中随机生成某种敌人。
这道题就是一个小巧但非常有代表性的例子。
我们要做的事情是:
给定一个整数数组 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 Foundationclass 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 Foundationclass Solution2 {private var nums: [Int]init(_ nums: [Int]) {self.nums = nums}func pick(_ target: Int) -> Int {var count = 0var result = -1for (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 或 4let solution2 = Solution2([1, 2, 3, 3, 3])
print(solution2.pick(3)) // 随机返回 2、3 或 4
可能输出结果:
4
0
3
2
你可以多运行几次,会发现随机返回的索引大致是均匀分布的。
这就说明算法的概率公平性是正确的。
时间复杂度
-
哈希表法:
- 初始化:O(n)
- 查询:O(1)
-
蓄水池抽样法:
- 初始化:O(1)
- 查询:O(n)
所以:
- 如果
pick
调用次数较多,推荐哈希表法; - 如果内存敏感,
nums
特别大,推荐蓄水池抽样。
空间复杂度
- 哈希表法: O(n),因为要保存每个数字的所有索引;
- 蓄水池抽样法: O(1),只用常量空间存储变量。
总结
这道题是一个非常有代表性的“概率算法”练习,它不仅考查你对随机性的理解,也考查你对空间-时间权衡的把握。
两种方法的对比总结如下:
方法 | 思路 | 查询效率 | 空间占用 | 适用场景 |
---|---|---|---|---|
哈希表法 | 预存索引列表 | O(1) | O(n) | pick 调用次数多 |
蓄水池抽样法 | 动态随机抽样 | O(n) | O(1) | 内存受限或数据流 |
如果你在做系统开发或分布式任务调度时,这个思路同样适用,比如:
- 从大量服务器节点中随机选择一个符合状态的;
- 从日志流里随机采样一条;
- 从用户群体中随机挑选符合条件的一人。