资讯动态

数组中的第K个最大元素:堆与快速选择深度剖析

发布时间:2026/9/12 5:36:46 来源:尧图企业网站定制
面试官递过来这道题的时候很多人第一反应是排序之后取第 k 个不就行了如果你真的当场只写一个 sort对方一般会追问一句“还能再快一点吗”。数组中的第K个最大元素力扣原题编号 215也是热题 100 里的常客。它号称中等难度却是理解两个高频算法思想的绝佳入口一个是堆另一个是快速选择。这篇文章不打算只贴代码而是想从“为什么用小顶堆而不是大顶堆”“为什么快速选择能到 O(n)”这些角度讲清楚再给出可以放心照抄的实现最后补一些我在实际刷题和面试复盘里踩过的坑。无论你是准备面试还是工作中遇到 TopK 类需求这篇应该都能帮上忙。1. 题目到底在考什么两个隐藏条件与“第K大”的下标换算1.1 题干里的两个隐藏条件原题描述很简短给定整数数组nums和整数k请返回数组中第k个最大的元素。看起来一句话就能看懂但有两个条件经常被忽略。第一个条件题目说的是“第 k 个最大”不是“第 k 个不同”。这意味着如果数组里有重复元素重复出现的值要重复计入排名。很多人在做示例 2 的时候会愣一下[3,2,3,1,2,4,5,5,6]、k 4答案为什么是 4因为按从大到小排序后是[6,5,5,4,3,3,2,2,1]第四个数是 4。这里的 5 出现了两次分别排第 2、第 3 位3 也出现了两次排第 5、第 6 位。如果你按“不同元素”去重数组会变成[6,5,4,3,2,1]第 4 大就变成 3直接做错。第二个条件数组本身是无序的而且题目不会保证数组规模小。力扣的限制是1 k nums.length 10^5。当数据量到十万这个级别时任何一个 O(n²) 的解法都会明显吃力O(nlogn) 能过但不够漂亮真正的考点是能不能把复杂度压到 O(nlogk) 甚至期望 O(n)。1.2 用示例 2 理解“重复元素也要计数”我们来手动过一遍示例 2因为它是很多人第一次提交错在哪里的最佳教材。数组是[3,2,3,1,2,4,5,5,6]k 4。先按从大到小排序6, 5, 5, 4, 3, 3, 2, 2, 1排第 1 的是 6排第 2 的是第一个 5排第 3 的是第二个 5排第 4 的是 4。所以答案是 4。如果换成“去重后取第 k 大”那答案会变成 3。出题人专门设计这个示例就是想提醒你不要把题目理解成集合去重。这个点在一些变种题里尤其容易踩比如“前 K 个高频元素”那里确实要按元素去重后再排名次但本题不是。1.3 排序视角下的下标换算表很多人写这道题最痛的一步不是算法本身而是下标换算。因为“第 k 大”和“升序排序后第几个位置”之间隔着一层转换。假设数组长度是 n第 k 大就是升序排序后从 0 开始计数、下标为n - k的那个元素。为了扎实可以记住下面这张表概念升序排序后的下标示例nums [3,2,1,5,6,4]n6第 1 大最大值n - 16下标 5第 2 大n - 25下标 4第 k 大n - k若 k2答案为 5第 n 大最小值01下标 0这个小表格建议刻在脑子里。后面快速选择算法里我们要求的就是“下标为 n-k 的元素”而堆解法里我们维护的堆顶也是第 k 大两个思路最终面对的是同一个目标。把这一层想清楚代码基本不会写错主逻辑。2. 排序解法为什么只是“第一版答案”而不是“最终答案”2.1 一行排序就完事的代码先说最直白的方案。Python 和 Java 各一行核心代码class Solution: def findKthLargest(self, nums: List[int], k: int) - int: nums.sort() return nums[-k]class Solution { public int findKthLargest(int[] nums, int k) { Arrays.sort(nums); return nums[nums.length - k]; } }Python 的nums[-k]就是倒着数第 k 个也就是升序排序后下标n-k的元素。Java 写成nums[nums.length - k]意思一样。这个解法在力扣上能通过因为数据规模不算变态内置排序的常数又小。如果你在笔试里时间紧张写这个不至于超时。但这不是一个“好”的解法原因在下面的复杂度分析。2.2 O(nlogn) 为什么会被追问优化排序的时间复杂度是 O(nlogn)空间复杂度取决于具体排序实现Python 的 TimSort 是 O(n) 辅助空间Java 对原始类型的双轴快速排序是 O(logn) 递归栈空间。关键问题是我们真的需要把所有元素排好吗题目只问第 k 大并不关心第 k 大之前的元素内部顺序也不关心第 k 大之后的元素谁大谁小。排序把整个数组整理得井井有条但我们只需要一个位置。这是一种“过度回答”。面试官追问“能不能优化”本质是看你能不能意识到这种过度。举一个极端情况k 1也就是找最大值。这时候任何人都会说一次遍历 O(n) 搞定不会有人先排序再取最后一个。同理k 很小的时候比如从十万个数里取第 10 大排序仍然做了很多无用功。所以这个问题的优化方向是找到一种方法让计算量尽量只和“我们需要的那部分”相关。2.3 排序方案在什么时候“够用”作为一个从业者我得说句实话在生产环境里80% 的情况下直接排序就够了。原因很简单真实业务里数组规模通常不大而语言内置的排序经过重度优化常数极小。一个 O(nlogn) 的排序在 n 10^5 时运行时间也就是几十毫秒完全不是瓶颈。如果团队里的同事为了“优化”去手写一个快速选择结果分区写得有 bug那才是灾难。但刷题和面试是另一套逻辑。面试官希望看到你掌握两种标准优化基于堆的 TopK 维护以及基于快速排序思想的选择算法。这两种思路也是下一节起的主角。3. 堆解法容量为K的小顶堆为什么是TopK问题的默认姿势3.1 先决定方向要找第K大为什么用小顶堆而不是大顶堆这是这道题最容易被问倒的点。如果你要找第 K 大正确的做法是维护一个容量为 K 的小顶堆。注意是小顶堆不是大顶堆。我用一句话解释背后的逻辑小顶堆的堆顶是堆里最小的元素也就是当前所有候选者里“最不配待在前 K 名”的那个。我们用这个堆顶当作门槛——如果新元素比门槛大就把门槛淘汰让新元素进场如果新元素比门槛小它连前 K 名都进不了直接忽略。对照一个生活里的场景假设你在选拔班级前 10 名现在手上已经有 10 个人的花名册花名册上分数最低的是第 10 名。来了一个新同学如果他的分数比花名册里最低的那个高就把最低的划掉让他顶上来如果比最低的还低那这个新同学显然进不了前十。反过来如果用了大顶堆堆顶是堆里最大的元素你根本不知道应该淘汰谁。你总不能把班级第一名淘汰掉还指望这个堆里剩下的人是全班前 10 名吧这就是为什么“找第 K 大用小顶堆找第 K 小用大顶堆”是 TopK 问题的铁律。3.2 容量为 K 的小顶堆Python heapq 实现Python 的heapq模块默认是小顶堆正好符合需求。import heapq def findKthLargest(nums, k): heap [] for num in nums: if len(heap) k: heapq.heappush(heap, num) else: # 等价于如果 num 大于堆顶就弹出堆顶压入 num heapq.heappushpop(heap, num) return heap[0]核心是heappushpop(heap, num)这个函数。它的行为是先把num压入堆再弹出堆中最小的元素。当堆的容量是 k 时这一步可以让堆保持大小为 k并且始终存着扫描过的元素里最大的 k 个。堆顶自然就是这 k 个里最小的那个也就是整个数组的第 k 大。如果你更习惯显式写法可以这样for num in nums: if len(heap) k: heapq.heappush(heap, num) elif num heap[0]: heapq.heapreplace(heap, num)注意heapreplace和heappushpop的区别。heapreplace是先弹出堆顶再压入新元素更适合“确定新元素比堆顶大”的场景heappushpop是先压入再弹出即使新元素比堆顶小堆结构也是安全的。两种写法在结果上没有差别选一种顺手且不容易记混的即可。另外送一个 Python 快得多的做法return heapq.nlargest(k, nums)[-1]。它内部也是用堆实现的但如果你想在面试里展示对原理的理解还是手写比较有说服力。3.3 Java PriorityQueue 的写法与默认排序Java 的PriorityQueue默认是小顶堆所以代码和 Python 非常像public int findKthLargest(int[] nums, int k) { PriorityQueueInteger heap new PriorityQueue(k); for (int num : nums) { if (heap.size() k) { heap.offer(num); } else if (num heap.peek()) { heap.poll(); heap.offer(num); } } return heap.peek(); }这里有个新手常踩的坑Java 的PriorityQueue默认是小顶堆不需要自定义比较器只有当你需要大顶堆时才要写成new PriorityQueue(Comparator.reverseOrder())也就是说找第 K 大用的是默认小顶堆很多人在这一步绕晕反而画蛇添足地传了一个reverseOrder()把堆变成了大顶堆结果堆顶永远是最大的那个逻辑全乱。3.4 复杂度和大数据场景的工程价值堆解法的时间复杂度是 O(nlogk)。每次heappush或heappop都是 O(logk)因为堆的规模是 k一共处理 n 个元素所以总复杂度是 O(nlogk)。空间复杂度 O(k)也就是堆本身占用的大小。当 k 很小的时候这个复杂度几乎可以看成 O(n)。比如 k 10logk ≈ 3.32实际就是每个元素做几次常数级的堆调整。当 k 接近 n 时O(nlogk) 就退化到接近 O(nlogn)这时候堆方案反而没有优势。这才是堆方案最值钱的地方它不要求数据一次性全部加载进内存而是可以流式处理。想象一个真实场景日志系统每分钟产生几百万条访问记录你想实时维护访问量最高的 100 个用户。如果先把所有记录存下来再排序内存早就爆了。但用一个容量为 100 的小顶堆每来一条记录就尝试入堆内存占用永远是 O(100)还能实时给出前 100 名单。这就是 MapReduce / Flink 等框架里做局部 TopK 的底层思路也是“工程上的堆”和“内存模型里的堆”经常被一起提到的原因。顺带一提这里说的堆是数据结构和 JVM 运行时那个“堆内存”“栈内存”完全是两码事。面试时如果被追问“堆和栈的区别”一定要先确认对方问的是数据结构还是内存模型否则容易答非所问。4. 快速选择把快排的Partition变成“定向搜索”期望O(n)的真相4.1 从快排到快选只做需要的“那一半”快速排序每次选择一个 pivot通过一次 partition 操作把数组分成“小于等于 pivot”和“大于 pivot”两部分pivot 落到它最终应该在的位置。快速选择就是从这里“偷懒”的既然 partition 之后 pivot 的下标已经固定而这个下标就是该元素在有序数组中的最终位置那我们只需要判断目标下标在 pivot 左边还是右边然后只处理那一边另一边完全不管。这就是为什么快速选择的期望复杂度是 O(n) 而不是 O(nlogn)。每次 partition 需要扫描的区间长度都在快速缩小第一次扫描 n 个元素下一次大概扫描 n/2 个再下一次 n/4 个……加起来是等比数列n n/2 n/4 ... 2n也就是 O(n)。对比一下快速排序每次 partition 之后两半边都要继续递归所以总工作量是 O(nlogn)快速选择只走一条分支省掉了一大半计算。4.2 第K大和第几小target n - k 的换算逻辑前面说过第 k 大对应升序排序后下标n - k。在快速选择里目标就是让一个元素落到target n - k这个位置。举个例子nums [3,2,1,5,6,4]n 6k 2那么target 4。第 2 大的元素是 5它在升序数组[1,2,3,4,5,6]中的下标确实是 4。这个位置就是我们的“定向目标”。如果你把题目换成“第 k 小”那目标下标直接就是k - 1。刷题时一定要先确认是第 k 大还是第 k 小不要把缓存里的模板无脑搬上去。4.3 双向扫描的写法很多Lomuto Partition 最好理解快速选择的 partition 写法有很多种Hoare 双边扫描不管实现还是理解都更费劲我建议先用 Lomuto 单边扫描面试时也更容易讲清楚。完整代码import random class Solution: def findKthLargest(self, nums: List[int], k: int) - int: target len(nums) - k def partition(l: int, r: int) - int: # 随机选择一个 pivot换到末尾 pivot_idx random.randint(l, r) nums[pivot_idx], nums[r] nums[r], nums[pivot_idx] pivot nums[r] i l for j in range(l, r): if nums[j] pivot: nums[i], nums[j] nums[j], nums[i] i 1 nums[i], nums[r] nums[r], nums[i] return i l, r 0, len(nums) - 1 while l r: pos partition(l, r) if pos target: return nums[pos] elif pos target: l pos 1 else: r pos - 1 return -1我用迭代写法而不是递归原因有两个一是避免递归深度过深二是面试时问到“如果用递归怎么写”你能给出迭代版会显得基本功更扎实。解释一下 Lomuto partition 里两个指针的含义。j是扫指针从l一直扫到r-1i是“小于等于 pivot 区域”的右边界。每当nums[j] pivot就把nums[j]交换到i的位置然后i后移一位。扫描结束后i指向的位置就是 pivot 应该待的位置交换nums[i]和nums[r]即可。这里用而不是是为了保证所有等于 pivot 的元素都被划到左半区避免在全是相等元素的数组里出现i无法前进的极端情况。4.4 期望O(n)、最坏O(n²)为什么随机化是刚需快速选择平均情况下是 O(n)但最坏情况会退化到 O(n²)。最坏情况的来源和快排一样如果每次选到的 pivot 都是当前区间的最小值或最大值那么 partition 之后左右极不均衡目标区间每次只缩小 1 个元素总计算量变成n (n-1) (n-2) ... O(n²)。什么时候会稳定触发这个最坏情况当数组已经排序而你每次固定取最后一个元素当 pivot 的时候。为了规避这个问题代码里必须随机选 pivot也就是先用random.randint(l, r)选一个下标再把它和nums[r]交换。随机化把最坏情况的概率摊开使得实际运行基本都贴合期望性能。如果你在面试里被问到“有没有办法让最坏情况也是 O(n)”可以提一下 BFPRT 算法中位数中的中位数它能保证最坏 O(n)但常数非常大实际工程里很少用。能说出来这个名字并解释“因为常数过大所以不值得用”已经能证明你研究过这部分了。5. 堆与快速选择怎么选从一张对比表到TopK衍生变体5.1 一张表看清三种方案的取舍方案时间复杂度空间复杂度是否原地支持流式数据推荐场景排序O(nlogn)O(1)~O(n)基本原地否数据量小、追求简单堆O(nlogk)O(k)需额外空间是数据量大、k 很小、流式场景快速选择期望 O(n)最坏 O(n²)O(1) 原地迭代版是否离线全量数据、追求最优时间这张表基本回答了一个核心问题在工程里到底选哪种如果数据已经全部在内存里且你只需要一次第 k 大查询快速选择的平均时间最优内存几乎零开销。如果数据是一个源源不断的流或者你需要在动态插入数据的过程中反复查询第 k 大的值堆是唯一合理的选择因为快速选择无法在流式数据上维护状态。5.2 衍生变体数据流第K大、最小K个数、前K高频词学会这道题相当于掌握了 TopK 问题族的引擎。现实中很多题只是给这个引擎换了层外壳。先看力扣 703“数据流中的第 K 大元素”。它要求不断往数组中添加新元素每次添加后返回当前第 k 大的元素。做法就是在初始化时维护一个容量为 k 的小顶堆后续每次add直接用前面的heappushpop逻辑即可。这就是同一个堆思路的“在线版本”。再看“最小的 K 个数”。对称地要用大顶堆维护容量为 k堆顶是当前 k 个数里最大的每次遇到比堆顶小的数就把堆顶淘汰、新数入堆。最终堆里就是全局最小的 k 个。还有“前 K 个高频元素”。先遍历整个数组用哈希表统计每个元素出现次数然后把(次数, 元素)二元组放进容量为 k 的小顶堆按次数排序。这里的堆顶依然是“淘汰线”只是比较的对象从元素大小换成了频次大小。很多同学会觉得“前 K 个高频元素”是道难题其实拆开看就是两个步骤哈希表计数 本题的堆模板。5.3 工程选型离线全量 vs 在线流式回到真实业务去看这三种方案选型就特别清晰。离线全量场景比如每天凌晨对昨天的日志跑一次 Top100 报表数据已经落在 HDFS 或数据库里。这时候最快的做法是直接ORDER BY xxx LIMIT 100数据库内核会自动选择排序或堆算法来执行。如果你想在自己代码里实现那用快速选择会更稳因为它无需额外内存且平均耗时最少。在线流式场景比如监控系统要实时展示当前 QPS 最高的 10 个接口每秒钟都有新数据进来。你不能每次来一条数据就把全部历史重新排序而堆可以做到 O(logk) 的增量更新。这也是为什么 Redis 的 ZSET、Flink 的 TopN 算子内部普遍采用类似堆的思路。工程里还有一种常见需求从单机 TopK 扩展到分布式 TopK。做法通常是每个节点先用自己的堆算出局部 TopK然后汇总到中心节点再用一个堆算出全局 TopK。局部 TopK 这一步用的就是本题的堆模板。5.4 中位数等更复杂的堆应用如果把堆的应用再往前推一步就到了“对顶堆”技巧一个大顶堆维护较小的一半一个小顶堆维护较大的一半两个堆的堆顶拼在一起就能得到中位数。力扣 295“数据流的中位数”就是最好例子。每次插入新数时先和大顶堆堆顶比较决定进哪个堆再通过调整两个堆的大小差不超过 1保持结构稳定。中位数要么是大顶堆堆顶要么是两个堆顶的平均数。这个技巧的本质就是“用两个容量不同的 TopK 堆拼接出让中间位置可见的结构”。理解了单堆 TopK再看对顶堆会顺畅很多。6. 实战里最容易翻车的细节随机化、重复元素、下标和边界6.1 随机化Pivot书写顺序很关键如果你把随机化写错顺序代码的行为会非常诡异。常见错误是随机选了下标但忘了把 pivot 先交换到最右端然后在 Lomuto partition 里仍然拿nums[r]当 pivot。这样你随机选的索引完全没参与分区随机化形同虚设。正确顺序是pivot_idx random.randint(l, r) nums[pivot_idx], nums[r] nums[r], nums[pivot_idx] pivot nums[r]先把随机选中的元素换到最后再用最后的元素做 pivot后续的交换逻辑才能统一。我见过有人在团队 Code Review 里提过“为什么快速选择也一定要随机化”的问题。答案是快速选择的最坏退化不是理论存在而是实际会发生。比如你已经知道测试数据是随机生成的固定取最后一个元素大概率没事但真实业务里“输入恰好近乎有序”是再常见不过的不随机化等于把性能交给运气。6.2 重复元素特别多的时候三路快选更安全前面的 Lomuto partition 用把所有等于 pivot 的元素归到左半区。当数组里重复元素很多比如一百万个数全是 1、2、3 循环等于 pivot 的元素会大量堆积在左半区分区仍然可能出现严重的不平衡导致性能劣化。这时候更稳妥的做法是三路快速选择把区间分成三块小于 pivot、等于 pivot、大于 pivot。def quick_select_three_way(l, r): if l r: return idx random.randint(l, r) pivot nums[idx] lt, i, gt l, l, r while i gt: if nums[i] pivot: nums[lt], nums[i] nums[i], nums[lt] lt 1 i 1 elif nums[i] pivot: nums[i], nums[gt] nums[gt], nums[i] gt - 1 else: i 1 if target lt: return quick_select_three_way(l, lt - 1) elif target gt: return quick_select_three_way(gt 1, r) else: return pivot三路快选最妙的地方在于如果 target 恰好落在“等于 pivot”这个区间可以直接返回。这意味着重复元素很多时一次 partition 就能结束战斗性能非常稳定。这一点也是荷兰国旗问题的经典思路面试时如果能主动提一句“对于重复元素多的场景我会用三路快选而不是普通 Lomuto”印象分会明显不一样。6.3 面试怎么“讲”这道题推荐的表达链很多人以为面试考的是“会不会写”实际上考的是“能不能讲清楚权衡”。我建议按下面这条线组织表达第一层先说排序方案O(nlogn)能过但不够好。面试官通常会点头然后问“能不能优化”。第二层给出堆方案维护一个小顶堆容量是 k堆顶就是第 k 大。时间复杂度 O(nlogk)空间 O(k)。突出它适合流式数据和 k 很小的场景。第三层追问“能不能再快”再给快速选择利用 partition 每次只处理目标所在的一侧期望 O(n)原地操作空间 O(1)。但必须补充最坏 O(n²) 的问题以及为什么随机化 pivot 能解决。第四层如果面试官问“K 接近 n 怎么办”可以说堆方案退化成接近 O(nlogn)这时候快速选择更有优势如果数据是流式输入堆仍然是唯一选择。这条表达链本身就体现了你对问题本质的理解排序是全量有序堆是“淘汰赛”快选是“定向搜索”。6.4 自查清单最容易写错的四行代码刷完题之后对照下面这个清单检查一遍能节省大量 Debug 时间。第一target len(nums) - k不是k - 1。写成后者是把第 k 大当成第 k 小处理了。第二Lomuto partition 里i的初始值是l不是0。很多人习惯性写 0子区间一旦不是从 0 开始交换就全乱了。第三缩减区间时用l pos 1和r pos - 1不要写成l pos或r pos。否则当目标在 pivot 右侧时区间可能永远不会缩小死循环。第四堆解法里判断条件是len(heap) k还是heap.size() k不同语言的 API 名称不同但逻辑都是“堆还不够容量就先入堆满了再比较堆顶”。这个条件写反会导致堆里最后不是 k 个元素。我最初刷这道题时在第二点和第三点上各翻过一次车。当时定位了很久才发现是自己 partition 的边界没处理好后来习惯性写完后先用一个随机小数组打印中间过程用不了几秒钟就能暴露问题。刷算法题最怕的不是不熟练而是单步调试时看不出来逻辑错在哪。如果你也卡在类似的地方建议把 partition 单独抽出来测试输入[5,1,4,2,3]这类小数组手动模拟一遍交换过程很快就能找到问题。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价