前几天有个准备秋招的朋友发消息问我LeetCode 398这道题代码我背下来了但面试官让我现场证明为什么这样选是等概率的我当场卡住了。这大概是很多刷题人的真实状态——题目本身叫 Random Pick Index翻译过来就是随机数索引代码短到可以闭眼默写但里面藏的蓄水池抽样Reservoir Sampling思想才是面试官真正想挖的东西。这篇文章我准备把这道题从原理到代码、从证明到面试话术全部掰开讲透。内容包括为什么最直观的哈希表解法不是最优解、蓄水池抽样为什么能扫描一遍就等概率、Python 和 Java 的双语言实现细节以及面试里考官最常追问的几个变形和翻车点。无论你是刚开始刷题的新手还是准备系统过一遍高频题的老手这篇应该都能帮上忙。1. 题目到底在考什么先分清随机返回下标的三种解法差异1.1 先手写最直觉的答案398 的题意一句话就能说清给定一个整数数组 nums可能包含大量重复元素。要实现一个 pick(target) 方法每次调用时从数组中所有等于 target 的下标里等概率返回任意一个。绝大多数人第一反应是预处理用哈希表把每个值出现过的下标都收集起来pick 的时候在对应列表里随机选一个。这个方案没有任何错误代码也很直白import random from collections import defaultdict class Solution: def __init__(self, nums): self.pos defaultdict(list) for i, num in enumerate(nums): self.pos[num].append(i) def pick(self, target): return random.choice(self.pos[target])复杂度也很清楚初始化要做一遍全量遍历时间复杂度 O(n)空间上要维护每个值的下标列表最坏情况是数组里只有一个数、出现了 n 次那就要存 n 个下标空间复杂度 O(n)。查询时是 O(1) 时间直接在列表里随机挑一个。这个答案能拿满分吗在力扣上能通过在面试里可能只能算及格。因为题目出现在蓄水池抽样这个标签下考官想听的解法不是哈希表而是空间 O(1)、只扫描一遍的蓄水池思路。1.2 但如果数据是流呢哈希表方案的空间隐患不妨顺着这个思路想一个场景假设 nums 不是一个固定下来的数组而是源源不断产生的数据流——比如后端实时打印的日志、用户不断产生的点击行为、传感器每秒上报的数据。你根本不知道总量是多少也没法提前把全部下标收集进哈希表。这时候哈希表方案就失效了它必须等数据全部到齐、落成数组之后才能建索引。而蓄水池抽样恰好天生就是为流式数据设计的数据来一个处理一个处理完就可以丢弃全程只保留一个或 K 个候选位置空间开销是常数级。所以对 398 这道题哈希表和蓄水池抽样都能做但背后的工程语义完全不同。这也是本文第一件要说清楚的事你写的每一行代码背后对应的是对数据是否可预知、是否可存储的假设。2. 蓄水池抽样为什么扫描一遍就能做到等概率2.1 核心思想逐个元素赌博式替换先说结论。蓄水池抽样在本题的精简版是这样的逻辑从左到右扫描数组。每遇到一个等于 target 的下标就记为第 cnt 个目标下标。对第 cnt 个下标以 1/cnt 的概率把它选为新候选否则保留之前的候选。扫描结束后候选下标就是最终的返回值。翻译成人话就是看到第一个目标下标时反正只有它一个直接选它看到第二个时掷一枚平均分的骰子有 1/2 概率换到第二个、1/2 概率留着第一个看到第三个时有 1/3 概率换到第三个、2/3 概率在原来两个里保留一个。这个过程非常像公司在年会上用击鼓传花的方式抽奖——每个到场的人都有机会成为最后的获奖者但概率取决于他在队伍里的位置。这个机制最反直觉的地方是你明明只保留了最后一个赢家却要求所有出现过的下标都有相同的被选中概率。为什么每个下标最终的概率不是越靠后越大这是理解这道题的关键也是面试官最想听到的推导。2.2 数学归纳法一次讲透假设目标值在整个数组中共出现 m 次下标按扫描顺序记为第 1 个、第 2 个、……、第 m 个。算法结束后我们要证明任意第 j 个下标成为最终候选的概率都是 1/m。用数学归纳法更准确的描述是处理完前 i 个目标下标后当前候选恰好是其中任意一个的概率都是 1/i。当 i 1 时只有一个候选它被选中的概率是 1即 1/1成立。假设处理完前 i-1 个目标下标时前 i-1 个中的每个下标成为候选的概率都是 1/(i-1)。现在来了第 i 个目标下标。算法以 1/i 的概率用它替换旧候选因此第 i 个下标成为新候选的概率就是 1/i。对任意 j i它在第 i 轮存活下来的前提是上一轮它是候选概率 1/(i-1)并且这一轮没有被替换概率 1 - 1/i (i-1)/i。两者相乘恰好是 (1/(i-1)) × ((i-1)/i) 1/i。所以处理完第 i 个下标后前 i 个下标每个仍有完全相同的概率 1/i。当 i 走到 m每个目标下标成为最终答案的概率就是 1/m。等概率成立。这个证明里最重要的两个数字是 1/i 和 1 - 1/i。前者保证新来者得到它应得的一份概率后者保证旧候选们把概率均匀让渡出来。一个在拿一个在让分毫不差。2.3 换个角度再看连乘消元的直观理解如果觉得归纳法太抽象还可以换一种更算术的理解方式。假设一个下标是第 k 个被扫到的目标下标。它要成为最终答案需要发生以下事件第 k 轮它被选中之后每一轮都不被替换。第 k 轮选中它的概率是 1/k第 k1 轮不被替换的概率是 1 - 1/(k1) k/(k1)第 k2 轮不被替换的概率是 (k1)/(k2)……第 m 轮不被替换的概率是 (m-1)/m。把这些乘起来1/k × k/(k1) × (k1)/(k2) × ... × (m-1)/m中间项全部约光剩下 1/m。你看k 无论取 1 还是 m-1最终概率都是 1/m。这就是蓄水池抽样公平的本质每一项分子分母前后相消位置靠前的下标靠多活几轮补偿位置靠后的下标靠选中的概率高补偿一来一去正好扯平。3. 双语言实现与边界细节3.1 Python 实现与随机 API 的选择蓄水池抽样在 398 题上的 Python 代码非常短import random class Solution: def __init__(self, nums): self.nums nums def pick(self, target): cnt 0 res -1 for i, num in enumerate(self.nums): if num target: cnt 1 if random.randint(0, cnt - 1) 0: res i return res注意随机 API 的用法random.randint(0, cnt - 1)返回的是闭区间 [0, cnt-1] 内的整数它等于 0 的概率正好是 1/cnt。这里也可以写成random.randint(1, cnt) cnt概率同样是 1/cnt但按习惯我建议统一用判断 0 的写法语义更直观一边遍历一边抽签抽到 0 号签就换人。之所以用整数随机而不是浮点数random.random() 1.0 / cnt是为了避开浮点精度问题。cnt 很小的时候两者没差别但 cnt 巨大时浮点数比较的边界情况多少有点隐忧。能用整数就别用浮点这是写随机算法时一条很实用的经验。3.2 Java 实现与 Random 的等价写法Java 版本逻辑完全相同只是随机 API 的边界要格外小心import java.util.Random; class Solution { private int[] nums; private Random rand; public Solution(int[] nums) { this.nums nums; this.rand new Random(); } public int pick(int target) { int cnt 0; int res -1; for (int i 0; i nums.length; i) { if (nums[i] target) { cnt; if (rand.nextInt(cnt) 0) { res i; } } } return res; } }rand.nextInt(cnt)返回 [0, cnt) 范围内的整数也就是 0 到 cnt-1判断等于 0 正好是 1/cnt 的概率。很多翻车现场都发生在这里nextInt的参数是上界不是个数写成nextInt(cnt 1)后概率就错成了 1/(cnt1)整个蓄水池抽样就不再均匀了。3.3 容易出错的三个实现细节代码虽然短但实现里有几个坑是 LeetCode 评论区常年被讨论的整理成表格方便对照细节错误写法正确写法原因随机数判断rand.nextInt(cnt) cntrand.nextInt(cnt) 0nextInt(cnt) 永远不会返回 cnt计数位置先判断随机再递增 cnt先递增 cnt 再随机第一个目标下标必须用 1/1 概率选中提前返回遇到 target 就随机判断并立即返回必须完整遍历整个数组提前返回会破坏后续下标的被选概率第三点值得单独展开说。很多第一次写蓄水池的人会想反正后面遇到的每个下标都有概率替换前面那我提前返回岂不是省时间大错特错。提前返回意味着只在前缀范围内做抽样如果 target 在后面还有大量下标它们永远没机会被选中整体概率立刻失衡。蓄水池抽样的前提就是必须看完所有数据一次都不能偷懒。4. 从 398 到通用蓄水池K 个样本与流式数据场景4.1 通用版从保留 1 个样本到保留 K 个样本398 只是蓄水池抽样最朴素的 K1 特例。通用问题是这样的有一个未知长度的数据流要在只遍历一遍的情况下从中等概率抽出 K 个样本。做法也有一脉相承的逻辑前 K 个数据直接放入蓄水池。从第 i 个数据开始i K以 K/i 的概率决定这个数据是否入选。如果入选就在蓄水池中随机挑一个位置替换掉。下面是一个完整的 Python 实现直接看比背概念有用import random def reservoir_sampling(stream, k): reservoir [] for i, item in enumerate(stream): if i k: reservoir.append(item) else: j random.randint(0, i) if j k: reservoir[j] item return reservoir这里的核心是第 i 个元素有 K/i 的概率被抽进池子等价于random.randint(0, i) k而池子里原有的元素也有各自的机会被顶掉。把 K1 代入就会得到j random.randint(0, i)判断j 1也就是j 0正好和 398 的写法对上。所以说白了398 的randint(0, cnt - 1) 0就是这个通用版公式的特殊形态。4.2 398 的亲兄弟们力扣上和这道题同源的还有几个熟面孔LeetCode 382链表随机节点。给一个单链表要求等概率返回任意一个节点的值。链表长度未知且只能走一遍天然就是蓄水池抽样的 K1 场景。LeetCode 384打乱数组。这个用 Fisher-Yates 洗牌算法解决和蓄水池是一对表亲核心都是每个位置和后续某位置交换的均匀随机思想。LeetCode 470用 Rand7 实现 Rand10。题目本身是另一类随机算法题但面试中经常和 398 连着问考察的是对随机分布的理解。如果你刷完 398 之后把 382 顺手刷掉大概率会发现一个是数组形态的蓄水池、一个是链表形态的蓄水池换汤不换药。这比盲目做十道新题更划算。4.3 业务场景数据流采样在工程里的真实用法面试之外蓄水池抽样真的在工业界被广泛使用。举几个我实际接触过的例子日志抽样。每天产生的日志可能有几十亿条如果想把 1% 的样本喂给离线分析最简单的方式不是先算出总量再抽样而是每来一条日志就以 1% 的概率决定留不留。系统不需要知道总量也不需要存储全部数据内存占用恒定。AB 实验流量分配。在一个大流量系统里做实验经常需要把用户请求按照比例随机分到对照组和实验组。如果实验组流量配比是动态调整的蓄水池抽样的思路能帮你在不知道总请求数的情况下保持相对均匀。推荐系统随机探索。给用户生成候选集合时有时需要从海量候选中均匀捞出几个作为多样性探索样本候选集合是实时算出来的根本不知道它有多大这时候蓄水池抽样就是很自然的解法。如果读者在做后端或者数据相关的工作下次遇到从一个大集合里均匀抽几个但集合大小不可知的需求可以下意识想想这道题的解法。它就是教科书和数据工程之间最短的连接点。5. 面试现场这道题的三个追问与常见翻车点5.1 追问一既然可以先数再随机为什么非用蓄水池这是面试官最爱抛出的挑战型问题很多候选人一听到就慌了。实际上对 398 的静态数组而言先统计一共有多少个 target 下标再随机选第几个并返回完全可以做到 O(n) 时间、O(1) 空间def pick(self, target): cnt sum(1 for x in self.nums if x target) k random.randint(0, cnt - 1) for i, x in enumerate(self.nums): if x target: if k 0: return i k - 1这方案错了吗没错。它的问题是必须完整遍历两次第一次数个数第二次定位下标。如果数据源是一个只能读取一次的流第一次遍历结束后数据就没了第二次遍历根本无从谈起。蓄水池抽样的核心优势恰恰是单趟扫描一边读一边维护候选读完即出结果。所以正确的面试回答姿势不是喷先数再随机的方案烂而是诚实地承认在静态数组场景下两种方案复杂度相当但如果把场景换成未知长度、单次遍历的数据流蓄水池就是唯一可用的方案。能把这道题辨析到这个深度比单纯背出一个蓄水池模板要加分得多。5.2 追问二哈希表预处理查询是 O(1)蓄水池是 O(n)不是退化了吗确实单次 pick 的时间复杂度蓄水池是 O(n)哈希表是 O(1)。但如果反复调用 pick两种情况的表现要分开算。哈希表方案预处理 O(n) 的时间和空间之后的每次 pick 都是 O(1)多次调用非常快代价是内存随着数据总量线性增长。蓄水池方案无需预处理空间 O(1)但每次 pick 都要重新扫描整个数组时间 O(n)。从复杂度角度看这是一个典型的空间换时间和时间换空间的权衡。在面试里考官往往自己想得到的也是这个权衡过程。实际工作中怎么选取决于数据规模数组只有几千长度但 pick 调用百万次哈希表明显更优数组大到内存装不下、或者本身就是流式数据蓄水池几乎是唯一解。5.3 追问三如果数组不变、pick 被频繁调用能不能优化蓄水池这题在力扣的讨论区里被称为隐藏的进阶版。既然蓄水池每次都要 O(n) 扫描而目标数组根本不变有没有办法做到第一次调用建立索引后续调用直接随机当然有。思路就是折中对每个不同的值只存它的出现次数不存全部下标。pick 时先随机决定要第几个出现的下标再从数组里定位。这样空间仍然远小于存下所有下标的哈希表但时间复杂度降不下来——定位还是要扫数组。如果真的是高频查询场景老老实实回到哈希表预处理反而更明智。LeetCode 的测试数据通常不会在这一点上卡你但面试中主动抛出这个权衡讨论会显得你不是在背模板而是真的理解每种方案的使用条件。5.4 现场回答的参考话术最后给一套可以直接用在面试里的回答逻辑大约两分钟说完整这道题我会先用蓄水池抽样的思路做。它的核心是遍历数组维护一个候选下标每当遇到第 cnt 个 target 下标就以 1/cnt 的概率替换候选下标。因为乘以 1/cnt 之后前面的候选被保留的概率是 (cnt-1)/cnt和新的 1/cnt 加起来正好让每个前 cnt 个下标概率均等。这里的关键是只能保存一个候选索引空间 O(1)成本是每次 pick 要完整扫描数组 O(n)。如果高频调用 pick 且数组能全部放入内存我会改用哈希表预处理所有下标用空间换时间。如果数据是流式的、长度未知蓄水池就是唯一的选择。这段话把原理、证明、复杂度和工程取舍全部带到了。哪怕部分细节说得不够严谨面试官也能听出你是系统学过这道题而不是只背了randint(0, cnt - 1) 0这一行。一点个人体会398 这道题我是先被面试官问住、之后才彻底搞懂的。当时我能写出代码却讲不清楚为什么1/cnt这个替换概率能保证全局公平直到自己动手做了连乘消元的推导才算真正拥有了这个算法。如果让我给后来人一个建议刷这道题的时候不要急着把答案背下来先给自己十分钟用笔把第 k 个下标最终成为答案的概率算一遍。算通了这道题和它的兄弟们——382、384、甚至通用蓄水池——就都拿下了。算不通代码背得再熟换个马甲还是会露馅。