Coding Interview Prep · 课时

Python heapq 与大根堆技巧

使用 heapq.heappush/heappop，将值取负来模拟大根堆，并应用 heapq.nlargest/nsmallest 快速查询前 k 个元素。

Python heapq 与大根堆技巧 是 CoddyKit 上的免费 Coding Interview Prep 课时。这是第 3 节课，共 4 节。你可以在下方免费阅读本课时的完整内容 — 然后在浏览器中使用内置代码编辑器和全天候 AI 导师进行实践。这是 Coding Interview Prep 学习路径的一部分，你的进度在网页和 CoddyKit 应用中同步。 Coding Interview Prep 课程共包含 4 节课。

Python 的 heapq 模块概览

Python 的 heapq 模块提供了一个基于普通 Python 列表实现的最小堆。与专用的堆类不同，heapq 会直接在现有列表上执行原地操作。该模块提供的函数包括：使用 O(n) 构建堆的 heapify，使用 O(log n) 添加元素的 heappush，使用 O(log n) 移除最小元素的 heappop，以及用于合并操作以提升效率的 heappushpop / heapreplace。

import heapq

# heapq operates on plain Python lists
heap = []
heapq.heappush(heap, 5)
heapq.heappush(heap, 2)
heapq.heappush(heap, 8)
heapq.heappush(heap, 1)

print('Heap array:', heap)          # internal array (not sorted!)
print('Peek min:', heap[0])         # O(1) min access
print('Pop min:', heapq.heappop(heap))  # 1
print('Next min:', heap[0])         # 2

# heapify: turn any list into a heap in O(n)
data = [9, 4, 7, 1, 3, 6, 2]
heapq.heapify(data)
print('Heapified:', data, '| min:', data[0])

通过取反值实现最大堆

Python 的 heapq 只提供最小堆。要模拟最大堆，请在推入之前将所有值取反，并在弹出时再次取反。这是因为堆按照存储的值排序，而取反会反转排序顺序。请务必记住两边都要取反：推入前取反，弹出后取反。遗漏其中任何一步，都是面试中常见的错误。

import heapq

max_heap = []
for val in [5, 1, 8, 3, 9, 2]:
    heapq.heappush(max_heap, -val)  # negate on push

print('Max-heap internal:', max_heap)  # all negated

# Pop in descending order:
results = []
while max_heap:
    results.append(-heapq.heappop(max_heap))  # negate on pop
print('Sorted descending:', results)  # [9, 8, 5, 3, 2, 1]

# Common pattern: top-k largest
data = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5]
k = 3
heap = []
for x in data:
    heapq.heappush(heap, -x)
print('Top', k, ':', [-heapq.heappop(heap) for _ in range(k)])

heapq.nlargest 和 nsmallest

heapq.nlargest(k, iterable) 和 heapq.nsmallest(k, iterable) 会返回最大的 k 个或最小的 k 个元素。它们的复杂度为 O(n log k)；当 k 远小于 n 时，比完整排序的 O(n log n) 更高效。它们在内部使用大小为 k 的堆。当 k 接近 n 时，Python 会退回到完整排序。对于不需要维护持久堆的一次性前 k 个查询，请使用这些函数。

import heapq

data = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 7]

# Top 3 largest:
print(heapq.nlargest(3, data))   # [9, 8, 7]
# Top 3 smallest:
print(heapq.nsmallest(3, data))  # [1, 1, 2]

# With a key function:
words = ['banana', 'apple', 'cherry', 'date', 'elderberry']
print(heapq.nlargest(2, words, key=len))   # ['elderberry', 'banana']
print(heapq.nsmallest(2, words, key=len))  # ['date', 'apple']

# Note: when k ~ n, use sorted() instead:
# sorted(data)[-k:]  or  sorted(data, reverse=True)[:k]

使用元组处理复杂键的堆

当堆元素需要自定义比较键时，可以将它们存储为元组 (priority, data)。Python 的 heapq 会逐个元素比较元组，因此会先比较优先级。如果优先级相同，它会比较第二个元素；如果数据不可比较，这可能会导致错误。最安全的模式是加入唯一计数器作为打破平局的依据，从而避免直接比较数据元素。

import heapq
import itertools

# Pattern: (priority, counter, item)
# Counter ensures unique tiebreaker, avoids comparing items
counter = itertools.count()
heap = []

def push_task(priority, task):
    heapq.heappush(heap, (priority, next(counter), task))

push_task(3, 'low priority task')
push_task(1, 'high priority task')
push_task(2, 'medium priority task')
push_task(1, 'another high priority')

while heap:
    pri, cnt, task = heapq.heappop(heap)
    print(f'P{pri}: {task}')
# Output in priority order: P1, P1, P2, P3

heapq.merge：合并有序可迭代对象

heapq.merge(*iterables) 会以惰性方式将多个有序可迭代对象合并为一个有序输出，而无需将所有数据加载到内存中。这相当于使用大小为 k 的最小堆执行 k 路归并，常用于外部排序算法。它返回一个迭代器，因此元素会逐个生成，非常适合大型数据集或流式处理场景。

import heapq

# Merge multiple sorted lists efficiently
sorted_lists = [
    [1, 5, 9],
    [2, 6, 8],
    [3, 4, 7]
]

# heapq.merge takes sorted iterables and returns a merged sorted iterator
merged = list(heapq.merge(*sorted_lists))
print('Merged:', merged)  # [1, 2, 3, 4, 5, 6, 7, 8, 9]

# The k-way merge manually (educational version):
def merge_k_sorted(lists):
    heap = []
    for i, lst in enumerate(lists):
        if lst:
            heapq.heappush(heap, (lst[0], i, 0))
    result = []
    while heap:
        val, list_idx, elem_idx = heapq.heappop(heap)
        result.append(val)
        if elem_idx + 1 < len(lists[list_idx]):
            heapq.heappush(heap, (lists[list_idx][elem_idx+1], list_idx, elem_idx+1))
    return result

print('Manual k-way:', merge_k_sorted(sorted_lists))

堆的延迟删除模式

当您需要从堆中删除任意元素，却不知道其索引时，可以使用延迟删除：在单独的集合中将元素标记为已删除，然后在弹出元素时跳过它们。其摊销复杂度为 O(log n)，并且避免了跟踪索引的复杂性。在包含重复条目的 dijkstra 算法和任务调度器模拟中，这是标准做法。

import heapq

class LazyHeap:
    def __init__(self):
        self._heap = []
        self._removed = set()

    def push(self, task):
        heapq.heappush(self._heap, task)

    def remove(self, task):
        self._removed.add(task)  # mark as removed

    def pop(self):
        while self._heap:
            task = heapq.heappop(self._heap)
            if task not in self._removed:
                return task
        return None

lh = LazyHeap()
for t in [5, 1, 8, 3, 2]:
    lh.push(t)
lh.remove(1)  # 'delete' 1 lazily
lh.remove(8)  # 'delete' 8 lazily
results = [lh.pop() for _ in range(3)]
print(results)  # [2, 3, 5] -- 1 and 8 skipped

数据流中的第 k 大元素

数据流中的第 k 大元素（LeetCode #703）维护一个大小为 k 的最小堆。堆顶始终是目前见过的第 k 大元素。当新数字到达时：将其推入堆中；如果堆的大小超过 k，则弹出最小元素。堆顶始终是第 k 大元素，因为堆中恰好有 k-1 个元素比它大。

import heapq

class KthLargest:
    def __init__(self, k, nums):
        self.k = k
        self.heap = []
        for num in nums:
            self.add(num)

    def add(self, val):
        heapq.heappush(self.heap, val)
        if len(self.heap) > self.k:
            heapq.heappop(self.heap)  # remove smallest
        return self.heap[0]  # kth largest = root of min-heap

# k=3, initial=[4,5,8,2]
kl = KthLargest(3, [4, 5, 8, 2])
print(kl.add(3))   # 4 (top 3: 8,5,4 -- kth=4)
print(kl.add(5))   # 5 (top 3: 8,5,5 -- kth=5)
print(kl.add(10))  # 5 (top 3: 10,8,5 -- kth=5)
print(kl.add(9))   # 8 (top 3: 10,9,8 -- kth=8)

查找和最小的 k 对

查找和最小的 k 对（LeetCode #373）使用最小堆按顺序生成元素对。首先，为每个 j 将所有元素对 (nums1[0], nums2[j]) 放入堆中。弹出最小元素后，对于弹出的元素对 (nums1[i], nums2[j])，将 (nums1[i+1], nums2[j]) 推入堆中；这是来自同一 nums2 列的下一个候选元素。这是使用堆按顺序生成元素对或乘积时的常见模式。

import heapq

def k_smallest_pairs(nums1, nums2, k):
    if not nums1 or not nums2:
        return []
    heap = []
    # Initialize with pairs (nums1[0], nums2[j])
    for j in range(min(k, len(nums2))):
        heapq.heappush(heap, (nums1[0] + nums2[j], 0, j))
    result = []
    while heap and len(result) < k:
        total, i, j = heapq.heappop(heap)
        result.append([nums1[i], nums2[j]])
        if i + 1 < len(nums1):
            heapq.heappush(heap, (nums1[i+1] + nums2[j], i+1, j))
    return result

print(k_smallest_pairs([1,7,11], [2,4,6], 3))
# [[1,2], [1,4], [1,6]]

使用最大堆的任务调度器

任务调度器（LeetCode #621）要求计算调度 n 个任务所需的最短时间，并规定相同任务之间必须间隔 n 个时间单位的冷却时间。请使用存储任务频率的最大堆：在每个时间步选择当前可用且频率最高的任务，将其计数减一，然后将其置于冷却状态。每个周期处理 k=n+1 个任务（或用空闲时间填充）。这种结合最大堆的贪心方法可以得到最优答案。

import heapq
from collections import Counter

def least_interval(tasks, n):
    freq = Counter(tasks)
    heap = [-f for f in freq.values()]  # max-heap (negated)
    heapq.heapify(heap)
    time = 0
    while heap:
        cycle = n + 1
        temp = []
        for _ in range(cycle):
            if heap:
                temp.append(heapq.heappop(heap))
        for f in temp:
            if f + 1 < 0:  # still tasks remaining
                heapq.heappush(heap, f + 1)
        # Add full cycle or remaining tasks if queue empty
        time += cycle if heap else len(temp)
    return time

print(least_interval(['A','A','A','B','B','B'], 2))  # 8
print(least_interval(['A','A','A','B','B','B'], 0))  # 6

dijkstra 算法中的堆

dijkstra 算法中的优先队列使用最小堆实现。存储元组 (distance, node)，始终先处理距离最近的未访问节点。当弹出节点的距离大于当前已知的最短路径距离时，说明这是延迟删除留下的过期条目，应跳过它。这样无需执行减小键值操作，就能保持实现简单，同时维持 O((V + E) log V) 的复杂度。

import heapq

def dijkstra(graph, start):
    dist = {node: float('inf') for node in graph}
    dist[start] = 0
    heap = [(0, start)]  # (distance, node)
    while heap:
        d, u = heapq.heappop(heap)
        if d > dist[u]:   # stale entry, skip
            continue
        for v, w in graph[u]:
            if dist[u] + w < dist[v]:
                dist[v] = dist[u] + w
                heapq.heappush(heap, (dist[v], v))
    return dist

graph = {
    'A': [('B', 4), ('C', 1)],
    'B': [('D', 1)],
    'C': [('B', 2), ('D', 5)],
    'D': []
}
print(dijkstra(graph, 'A'))  # {'A':0,'B':3,'C':1,'D':4}

使用最大堆重新排列字符串

重新排列字符串（LeetCode #767）要求重新排列字符串，使任意两个相邻字符都不相同。请使用由 (-frequency, char) 组成的最大堆。每一步都弹出频率最高的字符。如果前一个字符与频率最高的字符相同，则改为弹出频率第二高的字符。这种贪心方法会尽早放置限制最严格的字符。

import heapq
from collections import Counter

def reorganize_string(s):
    freq = Counter(s)
    heap = [(-f, c) for c, f in freq.items()]
    heapq.heapify(heap)
    result = []
    prev_freq, prev_char = 0, ''
    while heap:
        freq, char = heapq.heappop(heap)
        result.append(char)
        # Push back the previous character if still remaining
        if prev_freq < 0:
            heapq.heappush(heap, (prev_freq, prev_char))
        prev_freq, prev_char = freq + 1, char  # decrement freq (less negative)
    result_str = ''.join(result)
    # Verify no adjacent duplicates
    return result_str if len(result_str) == len(s) else ''

print(reorganize_string('aab'))   # 'aba'
print(reorganize_string('aaab'))  # '' (impossible)

快速检查

测试您对本课中数据结构与算法——编程面试准备相关概念的理解。

课程回顾

在本课中，您学习了：Python 的 heapq 模块及其 API，包括 heapify、heappush、heappop、nlargest、nsmallest 和 merge；通过取反值来模拟最大堆；以及包括前 k 个元素流式处理、数据流中的第 k 大元素、任务调度器和 dijkstra 在内的常见堆面试模式。接下来，我们将学习数据流中的中位数和 k 路归并。

常见问题解答

「Python heapq 与大根堆技巧」课时是免费的吗？

是的 — 「Python heapq 与大根堆技巧」的完整文本可在网页上免费阅读。要进行交互式练习（内置代码编辑器和全天候 AI 导师）并解锁 Coding Interview Prep 课程的其余内容，请升级到 CoddyKit PRO。 Coding Interview Prep 课程共包含 4 节课。

「Python heapq 与大根堆技巧」这节课中我会学到什么？

使用 heapq.heappush/heappop，将值取负来模拟大根堆，并应用 heapq.nlargest/nsmallest 快速查询前 k 个元素。你通过在浏览器中直接运行的动手代码来练习 Coding Interview Prep，全天候 AI 导师会在你学习这节课的过程中回答你的问题。

学习 Coding Interview Prep 需要有经验吗？

无需任何先前经验。CoddyKit 上的 Coding Interview Prep 课程适合初学者到高级学习者，你可以从这里开始或从头开始，按照自己的节奏学习。这是第 3 节课，共 4 节。

「Python heapq 与大根堆技巧」课时需要多长时间？

大多数 CoddyKit 课程大约需要 5–10 分钟。每节课都很精短且互动，所以你能稳步进步，并在网页和应用中从离开的地方继续。

我能在这节 Coding Interview Prep 课中编写并运行代码吗？

能。每节 Coding Interview Prep 课都包含内置代码编辑器，你可以在浏览器中直接编写并运行真实代码，并获得即时 AI 反馈 — 无需本地设置。

此课程中的所有课时

← 返回 Coding Interview Prep