Pattern 08 · Heap / Top-K

08 堆 / Top-K 只关心第 K 个时,不排序,养一个大小为 K 的堆

面试频率 高题目 6 道模板 2 个⭐认脸 1 道
§1 Mental Model 一句话心智模型

先记住这一句

只关心第 K 个的时候不要排序:养一个大小为 K 的堆,让不够格的自动掉出去。

方向最容易记反,记住这一句就够:求第 K 大,用大小为 K 的小顶堆。堆顶是这 K 个里最小的,也就是全局第 K 大;新来的数比堆顶大才有资格挤掉它。求第 K 小就整体取相反数,别去手写大顶堆——Python 3.13 及以前的 heapq 只有小顶堆(3.14 起新增 heappush_max 等,先确认面试环境的版本)。

复杂度是 O(n log K) 而不是排序的 O(n log n)。K 很小时差距巨大,而且它天然支持数据流:数字一个个来,你随时能回答,排序做不到这件事。

§2 Recognition Signals 识别信号

看到什么,想到本章

题面出现…想到
「第 K 大 / 第 K 小」且数组无序大小为 K 的堆(模板 A)
「前 K 个高频 / 最接近原点的 K 个点」模板 A,只换入堆的键
「数据流 / 不断插入,随时要回答」堆(排序做不到,因为数据没有终点)
「合并 K 个有序链表 / 数组」堆里放 K 个当前头,弹一个补一个
「动态中位数」对顶堆(模板 B)
「每次取出当前最大的,处理完再放回去」堆(621 任务调度、1046 碰撞石头)
分流条件 · 决策树同源

「前 K 个高频」要先用 Counter 数出频次(第 01 章)再用堆选 K 个——哈希负责统计,堆负责挑选,两步不要混。数组已排序时第 K 大就是一个下标,别掏堆;给的是 BST 时中序走到第 K 个即可(第 07 章)——「第 K」这三个字本身不构成用堆的理由,无序才是;「固定窗口内的最大值」要的是过期淘汰,那是单调队列(第 04 章)不是堆;「第 K 小的距离 / 乘积」这类答案是数值的,回到答案二分(第 05 章)。只要一次的第 K 大且允许平均 O(n),快速选择(quickselect)也是标准答案。

§3 Template 模板代码

可默写的那几行

■ 【骨架】 不变,闭着眼默写 ■ 【插槽】 随题变化,写题前先问自己填什么

模板 A · 大小为 K 的堆覆盖本章 80% 的题
import heapq

def kth_largest(nums, k):
    heap = []                               # 【骨架】小顶堆,堆顶是这 K 个里最小的那个
    for x in nums:                          # 【骨架】
        heapq.heappush(heap, x)             # 【插槽】入堆的键:值 / 频次 / 距离的相反数
        if len(heap) > k:                   # 【骨架】超员就把最弱的踢出去
            heapq.heappop(heap)             # 【骨架】
    return heap[0]                          # 【插槽】答案:第 K 大 / 整个 K 集合
默写口诀:入堆、超员、踢堆顶。换题只改一件事——入堆的键(值 / 频次 / 距离的相反数)。
模板 B · 对顶堆仅动态中位数需要,认脸级
import heapq

class MedianFinder:
    def __init__(self):
        self.lo = []                        # 【骨架】大顶堆(存相反数),装较小的一半
        self.hi = []                        # 【骨架】小顶堆,装较大的一半

    def add(self, x):
        heapq.heappush(self.lo, -x)                        # 【骨架】先无脑进左边
        heapq.heappush(self.hi, -heapq.heappop(self.lo))   # 【骨架】把左边最大的让给右边
        if len(self.hi) > len(self.lo):                    # 【插槽】平衡策略:允许谁多一个
            heapq.heappush(self.lo, -heapq.heappop(self.hi))

    def median(self):
        if len(self.lo) > len(self.hi):     # 【插槽】奇数个时中位数在较长的那边
            return float(-self.lo[0])
        return (-self.lo[0] + self.hi[0]) / 2
两个堆的不变式:lo 全体 ≤ hi 全体,且 len(lo) - len(hi) ∈ {0, 1}。「先进左边、再让给右边、右边多了再还回来」这三步能保证不变式恒成立。
§4 Problems 题目分层

从建直觉到扭模板

热身建直觉,easy 为主1 道
  • 模板 A 的类封装版,把「数据流」这件事具象化:构造时先灌一遍,之后每次 add 都是「入堆 + 超员就踢」。
标准模板直接套3 道
  • 模板 A 原题。面试标准答法是先说堆 O(n log K),再补一句「要平均 O(n) 就用快速选择」——两个都说得出才算完整。
  • Counter + 堆的组合。heapq.nlargest(k, cnt, key=cnt.get) 一行也行,但要能说出它内部就是这个堆。桶排序 O(n) 解法是加分项。
  • 入堆的键换成距离的相反数(求最近的 K 个 = 踢掉最远的)。别开平方,比较平方距离就够——省一次浮点运算,也避免精度问题。
变体模板要扭,坑的来源2 道
  • 堆解法是「每轮取出剩余最多的几种任务」,直观但啰嗦;面试更爱听 (maxCount - 1) * (n + 1) + 同频次任务数 这个数学公式。两种都准备,先讲公式再说堆。
  • ⭐认脸。模板 B 原题,说清两个堆的不变式即可。
§5 Pitfalls & Self-check 常见翻车点 + 自测

会写不等于不翻车

  1. 方向记反——求第 K 大用小顶堆。反过来(大顶堆存 K 个)踢掉的是最大的,答案全错。记法:堆顶是「守门员」,它守的是第 K 名这道门。
  2. 3.14 以前没有大顶堆——一律取相反数入堆,出堆再取回来,这在所有版本都能用。Python 3.14 起可用 heappush_max / heappop_max,但面试环境的 Python 未必 ≥ 3.14,先确认版本。对顶堆里 lo 存的是负值,-self.lo[0] 才是真正的最大值——这一步漏掉是 295 最常见的 bug。
  3. 元组入堆的次级比较——(freq, val) 里 freq 相同时会去比 val。如果 val 是不可比较的对象(如自定义类、字典),会直接抛 TypeError。塞一个序号当 tie-breaker。
  4. heapify 与逐个 push 的差别——heapq.heapify(list) 是 O(n),逐个 heappush 是 O(n log n)。已经拿到全量数据时用前者;数据流场景只能用后者。
  5. K 等于 n 时堆没意义——K 接近 n 时 O(n log K) 已退化到 O(n log n),直接排序更省事。面试里主动说这句,说明你知道自己在优化什么。

自测 · 30 秒识别模式

30 秒识别Q1 · 1046 Last Stone Weight:每次取出最重的两块石头相撞,剩下的差值放回去,求最后剩下的重量。你的第一反应是哪个模式?

看答案

堆。「每次取当前最大、处理完再放回去」是堆的标准句式。Python 里全部取负入小顶堆即可。排序做不到——每撞一次集合就变了。

30 秒识别Q2 · 692 Top K Frequent Words:返回出现频率前 K 高的单词,频率相同的按字典序排。你的第一反应是哪个模式?

看答案

Counter + 堆(模板 A)。频率相同按字典序,正是「元组入堆的次级比较」那条翻车点的实战:入堆键写成 (freq, 反向的字典序) 才能一次排对。

30 秒识别Q3 · 658 Find K Closest Elements:给一个已排序数组和目标值 x,找出最接近 x 的 k 个数。你的第一反应是哪个模式?

看答案

二分 / 双指针,不是本章——看到「k 个最接近」手就想去掏堆,但数组已排序,答案必然是一段连续区间,二分找左边界即可,O(log n)。堆要先把 n 个元素全扫一遍(O(n log k)),等于把题目给的「有序」这个最强条件原样扔掉。见 §2 分流条件。

每章自测都会混入恰好一道非本章题,防止「读哪章选哪章」的条件反射;题面也一律取自本章未收录的题,避免凭记忆而非凭信号作答。