LeetCode 398. 随机数索引

约 7 分钟
文章信息 8 个分类 4 个标签
分类
LeetCode CodeTop 中等 未解决 2-5 字节跳动 快手 拼多多
标签
哈希表 数学 随机化 水塘抽样
目录

题目描述

✅ 398. 随机数索引

题意分析

要设计一个类：构造时接收一个可能含重复元素的整数数组，之后支持 pick(target) 操作，返回 target 在数组中某个出现位置的下标。要求是——如果 target 出现了 $m$ 次，那么这 $m$ 个下标中每一个被返回的概率都必须是 $1/m$。

最重要的约束信号是「构造一次、查询多次」这个接口形状。数组一经传入就不再变化，而 pick 的调用次数可以很大，所以任何不随查询改变的信息都值得在构造期一次性算好。

第二个信号是题目明确保证 target 一定存在于数组中，因此不需要考虑空结果，也不需要为「找不到」设计返回值。

「等概率」这个词是判题的核心：返回一个合法下标很容易，难的是保证分布均匀。任何固定顺序、轮转或者偏向某一端的实现都会被判错。

边界包括：target 只出现一次（此时必然返回那唯一的下标）；数组里所有元素都等于 target（等价于在整个下标区间上均匀取样）；以及数组中存在大量不同值，预处理结构的规模会接近数组本身。

解法：哈希表预处理下标

核心思路

朴素写法是把工作全放在 pick 里：每次调用都完整扫一遍数组，把所有等于 target 的下标收集到一个列表，再从列表里随机取一个。这个做法的分布是正确的，问题出在代价上——单次查询 $O(n)$，$q$ 次查询就是 $O(nq)$，在 $n$ 与 $q$ 都上万时是亿级操作。

瓶颈非常具体：数组内容从构造之后就是只读的，可每次 pick 都要重新扫一遍并把扫描结果丢弃，同样的工作被重复了 $q$ 遍。

观察到「值到下标集合」的映射是完全静态的，就可以把它挪到构造函数里算一次：遍历数组，把每个下标追加到它对应数值的列表末尾。之后每次 pick 只需一次哈希查找拿到列表，再在列表长度范围内取一个均匀随机整数即可。

分布的正确性来自一条简单事实：若列表长度为 $m$，在 $[0, m)$ 上取均匀随机整数，每个位置被取到的概率都是 $1/m$，与它保存的具体下标无关。因此每个出现位置被返回的概率都是 $1/m$，恰好满足题目要求。

这里维持的不变量是：构造完成后，positions[v] 严格等于「v 在原数组中全部出现位置的升序列表」，并且此后不再被修改；pick 只读不写，任意两次调用互不影响，各自独立地均匀取样。

解题步骤

  • 在构造函数里建立一个「数值到下标列表」的哈希表。选哈希表而不是数组，是因为元素取值范围覆盖整个整型，无法按值开桶。
  • 顺序遍历 nums，把下标 i 追加到 nums[i] 对应的列表末尾。顺序追加自然保证列表升序，虽然本题不依赖顺序，但它让结构可预测、便于调试。
  • pick(target) 时先做一次哈希查找取出下标列表。题目保证 target 存在，所以不需要判空；若要写健壮版本，这里是唯一需要兜底的位置。
  • 在 $[0, m)$ 区间取一个均匀随机整数作为列表位置，其中 $m$ 是列表长度。上界必须是列表长度而不是数组长度，否则既会越界又会破坏分布。
  • 返回该位置保存的原数组下标。整个 pick 只做一次哈希查找和一次随机数生成，没有任何扫描。

以 nums = [1, 2, 3, 3, 3] 走一遍：构造阶段依次处理下标 0 到 4，positions 逐步变成 {1: [0]}、{1: [0], 2: [1]}、{1: [0], 2: [1], 3: [2]}、{..., 3: [2, 3]}，最终为 {1: [0], 2: [1], 3: [2, 3, 4]}。调用 pick(3)：取出列表 [2, 3, 4]，长度 $m = 3$，随机数落在 {0, 1, 2} 上各占 $1/3$，对应返回 2、3、4，恰好每个下标 $1/3$。调用 pick(1)：取出列表 [0]，长度为 1，随机数只能是 0，返回下标 0，概率为 1，符合「只出现一次时必然返回该位置」。

代码实现

class Solution {
    // 题目保证查询的 target 一定存在，因此取到下标列表后不需要额外处理空结果。
    private final Map<Integer, List<Integer>> positions = new HashMap<>();

    public Solution(int[] nums) {
        for (int i = 0; i < nums.length; i++) {
            positions.computeIfAbsent(nums[i], key -> new ArrayList<>()).add(i);
        }
    }

    public int pick(int target) {
        List<Integer> indices = positions.get(target);
        int idx = java.util.concurrent.ThreadLocalRandom.current().nextInt(indices.size());
        return indices.get(idx);
    }
}

type Solution struct {
    // 题目保证查询的 target 一定存在，因此取到下标列表后不需要额外处理空结果。
    positions map[int][]int
}

func Constructor(nums []int) Solution {
    positions := make(map[int][]int)
    for i, num := range nums {
        positions[num] = append(positions[num], i)
    }

    return Solution{positions: positions}
}

func (s *Solution) Pick(target int) int {
    indices := s.positions[target]
    idx := rand.Intn(len(indices))
    return indices[idx]
}

复杂度分析

  • 时间复杂度：构造函数 $O(n)$，每个下标只被追加一次；pick 期望 $O(1)$，代价是一次哈希查找加一次随机数生成，与 target 的出现次数无关。
  • 空间复杂度：$O(n)$，所有列表的元素总数恰好等于数组长度，再加上哈希表本身与不同值个数成正比的桶开销。

关键点总结

  • 「构造一次、查询多次」的接口形状是预处理换查询的强信号。判断依据是：这部分计算结果会不会随查询参数变化——不变的就该挪到构造期。
  • 等概率取样的论证要落在「均匀随机整数的取值范围」上：范围必须恰好是候选个数，多一个会越界，少一个会让末尾的候选永远取不到。
  • 空间换时间的代价必须能说清楚。这里用 $O(n)$ 额外空间换 $O(1)$ 查询；一旦面试官把空间限死在常数，方案就必须整体换掉。
  • 随机源应当复用同一个实例而不是每次新建，既避免种子相近导致的相关性，也省掉重复初始化的开销。
  • 面试视角：这题最常见的追问是「如果数组太大存不下、或者只能顺序读一遍怎么办」。答案是蓄水池抽样——扫描中遇到第 $k$ 个 target 时以 $1/k$ 的概率用它替换当前答案，最终第 $i$ 个被保留的概率是 $\frac{1}{i} \cdot \prod_{j=i+1}^{m}\left(1 - \frac{1}{j}\right) = \frac{1}{m}$。能把这个乘积推完，比写出代码更有说服力。
  • 面试视角：主动对比两套方案的取舍——哈希预处理是 $O(n)$ 空间、$O(1)$ 查询，蓄水池是 $O(1)$ 空间、$O(n)$ 查询，选哪个取决于查询次数与内存约束。展示「按场景选型」的思路是设计题的得分点。

易错点总结

  • 错误写法：随机数上界写成数组长度而不是列表长度。用例 nums = [1, 2, 3, 3, 3] 调用 pick(3) → 随机数可能取到 3 或 4，而列表长度只有 3，访问时下标越界抛异常。
  • 错误写法：构造时用覆盖式写入，只记录每个值最后出现的下标。用例 nums = [1, 3, 3] 调用 pick(3) → 恒定返回 2，下标 1 的概率是 0，等概率要求直接失败。
  • 错误写法：把收集下标的工作留在 pick 里现场做。用例 长度 $2 \times 10^4$ 的数组配 $10^4$ 次查询 → 总操作量到达 $2 \times 10^8$ 级别，直接超时。
  • 错误写法：用 random.nextInt() % m 取模代替带上界的随机。用例 nextInt() 返回负值 → 取模结果为负导致下标越界；即使先取绝对值，取模也会让前几个位置的概率偏高，分布不再均匀。
  • 错误写法：把随机结果缓存起来复用，只在第一次 pick 时抽一次。用例 连续两次 pick(3) → 返回同一个下标，违背「每次调用都独立等概率」的判题要求。
  • 错误写法：在每次 pick 内部新建随机数生成器实例。用例 同一毫秒内密集调用 → 若实现按当前时间播种，多次生成的种子相近，返回值高度相关，随机性被破坏。
  • 错误写法：为了「看起来公平」而按轮转顺序依次返回下标。用例 反复调用 pick(3) → 输出是 2、3、4、2、3、4 的固定序列，虽然覆盖全部位置但不是随机分布，判题会失败。
  • 错误写法：构造时只保存数组引用，把下标列表的构建推迟到首次查询，而数组在此期间被外部修改。用例 构造后外部把某个位置的值改掉 → 预处理表与实际数据不一致，返回的下标上根本不是 target。

相似题目

题目 难度 考察点
382. 链表随机节点 中等 长度未知且只能顺序访问，必须用蓄水池抽样在线维护候选
528. 按权重随机选择 中等 目标分布非均匀，用前缀和加二分把均匀随机数映射到带权区间
384. 打乱数组 中等 要求每一种排列等概率，考察 Fisher-Yates 原地洗牌的正确交换范围
380. O(1) 时间插入、删除和获取随机元素 中等 集合可变，需要变长数组配哈希表在支持等概率取样的同时做到常数删除