Python heapq模块实现优先队列与Top-K问题

在处理数据时，我们常遇到两种典型需求：一是需要一个能自动排序并快速取出最值的数据结构（优先队列）；二是从海量数据中快速找出前K大或前K小的元素（Top-K问题）。Python内置的heapq模块正是解决这类问题的利器。

本文将用最直接的方式，带你掌握用heapq实现优先队列与Top-K问题的全部核心操作。

第一步：理解heapq与“堆”

heapq模块提供了一个基于列表（list）的堆（heap）实现。你可以把“堆”想象成一个特殊的二叉树，它满足堆属性：对于一个小顶堆，任意父节点的值都小于或等于其子节点的值；反之，对于大顶堆，父节点的值大于或等于子节点的值。

heapq模块默认创建的是小顶堆。这意味着，通过heapq操作后的列表，其第一个元素（索引为0）永远是当前堆中的最小值。这正是实现优先队列的基础。

heapq模块中的所有操作，都是直接在传入的列表上进行修改，而不是返回一个新的列表。

第二步：掌握heapq核心函数

我们通过一个简单的例子来认识关键函数。

1. 导入模块。

   import heapq

2. 将列表“堆化”。使用heapq.heapify()函数，可以将一个普通列表就地转换为一个有效的堆。转换后，列表的顺序会发生变化，以满足堆属性。

   numbers = [5, 2, 8, 1, 9]
   heapq.heapify(numbers)
   print(numbers)  # 输出: [1, 2, 8, 5, 9] (列表顺序已改变，最小值1在开头)

3. 弹出最小值。使用heapq.heappop()函数，它会从堆中移除并返回最小的元素，同时自动维护堆结构，确保下一个最小值成为新的堆顶。

   smallest = heapq.heappop(numbers)
   print(smallest)  # 输出: 1
   print(numbers)    # 输出: [2, 5, 8, 9] (堆重新调整)

4. 插入新元素。使用heapq.heappush()函数，可以将一个新元素插入堆中，并保持堆属性。

   heapq.heappush(numbers, 3)
   print(numbers)  # 输出: [2, 3, 8, 5, 9] (新元素3被插入，堆重新调整)

5. 查看最小值但不弹出。直接访问列表的第一个元素 heap[0]。

   current_smallest = numbers[0]
   print(current_smallest)  # 输出: 2

通过组合heappush和heappop，我们就实现了一个优先队列：heappush负责按优先级（值的大小）排队，heappop总能取出优先级最高（值最小）的元素。

第三步：实现优先队列

一个完整的优先队列类通常包含以下方法。我们使用heapq来实现它。

import heapq

class PriorityQueue:
    def __init__(self):
        self._queue = []
        self._index = 0  # 用于在值相同时，按插入顺序排序

    def push(self, item, priority):
        # heapq是按元组比较的，第一个元素是优先级。
        # 加入index可以确保即使优先级相同，也能按插入顺序弹出
        heapq.heappush(self._queue, (priority, self._index, item))
        self._index += 1

    def pop(self):
        # heappop返回元组 (priority, index, item)，我们只取item
        return heapq.heappop(self._queue)[-1]

    def is_empty(self):
        return not self._queue

# 使用示例
pq = PriorityQueue()
pq.push('任务A', 2)
pq.push('任务B', 1)  # 更高优先级（数值更小）
pq.push('任务C', 2)

print(pq.pop())  # 输出: '任务B' (优先级最高)
print(pq.pop())  # 输出: '任务A' (优先级相同，但先被插入)
print(pq.pop())  # 输出: '任务C'

在这个实现中，我们将一个元组 (priority, index, item) 推入堆。Python元组的比较规则是：先比较第一个元素，如果相等再比较第二个，以此类推。_index保证了在优先级相同时，先入队的元素先出（FIFO）。

第四步：解决Top-K问题

Top-K问题是优先队列的典型应用场景。有两种思路：

思路一：使用小顶堆维护一个大小为K的“候选池”（适用于求前K大）。

• 维护一个包含K个元素的小顶堆。
• 遍历所有数据，如果当前元素大于堆顶，则弹出堆顶，并将当前元素压入堆。
• 遍历结束后，堆中剩下的K个元素就是前K大的元素。
• 复杂度分析：遍历N个元素，每次堆操作复杂度为$O(\log K)$，总复杂度为$O(N \log K)$。当$K \ll N$时，效率很高。

import heapq

def find_top_k_largest(nums, k):
    # 初始化一个大小为k的小顶堆，用前k个元素
    min_heap = nums[:k]
    heapq.heapify(min_heap)
    # 遍历剩余元素
    for num in nums[k:]:
        # 如果当前元素大于堆顶（当前堆中的最小值）
        if num > min_heap[0]:
            # 弹出堆顶（淘汰掉当前堆中最小的），压入新元素
            heapq.heapreplace(min_heap, num) # 比先pop再push更高效
    return min_heap

# 使用示例
data = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5]
top_3_largest = find_top_k_largest(data, 3)
print(top_3_largest)  # 输出: [5, 6, 9] (顺序不定，但确实是最大的三个数)

思路二：直接使用heapq.nlargest和heapq.nsmallest函数。
heapq模块提供了更便捷的函数，内部实现了上述优化逻辑。

import heapq

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

# 获取最大的3个元素
top_3_largest = heapq.nlargest(3, data)
print(top_3_largest)  # 输出: [9, 6, 5]

# 获取最小的3个元素
top_3_smallest = heapq.nsmallest(3, data)
print(top_3_smallest) # 输出: [1, 1, 2]

选择哪种？

• 如果只需要前K个结果，且数据流很大，用思路一的手动堆，可以严格控制内存为K。
• 如果数据已全部在列表中，且代码简洁性更重要，直接使用heapq.nlargest或heapq.nsmallest。

第五步：处理复杂数据结构的Top-K

实际问题中，我们往往需要根据对象的某个属性（如分数、价格、时间）进行排序。这时需要提供一个排序键函数。

假设我们有一组学生信息，需要找出分数最高的3个学生：

import heapq

students = [
    {'name': 'Alice', 'score': 85},
    {'name': 'Bob', 'score': 92},
    {'name': 'Charlie', 'score': 78},
    {'name': 'David', 'score': 95},
    {'name': 'Eve', 'score': 88}
]

# 方法1：使用nlargest，并指定key
top_3_by_score = heapq.nlargest(3, students, key=lambda s: s['score'])
print(top_3_by_score)
# 输出: [{'name': 'David', 'score': 95}, {'name': 'Bob', 'score': 92}, {'name': 'Eve', 'score': 88}]

# 方法2：手动实现（更灵活，可嵌入自定义逻辑）
def get_top_k_students(students, k):
    # 使用 (score, index, student_dict) 的元组
    indexed_scores = [(s['score'], i, s) for i, s in enumerate(students)]
    top_k_heap = indexed_scores[:k]
    heapq.heapify(top_k_heap)
    for item in indexed_scores[k:]:
        if item[0] > top_k_heap[0][0]:
            heapq.heapreplace(top_k_heap, item)
    # 从结果中提取出student_dict，并按分数降序排序（可选）
    top_k_students = [item[2] for item in top_k_heap]
    top_k_students.sort(key=lambda s: s['score'], reverse=True)
    return top_k_students

top_3_manual = get_top_k_students(students, 3)
print(top_3_manual)

当数据结构更复杂，或者需要更精细的控制（如去重、合并多个数据源）时，手动实现堆往往更合适。

通过以上步骤，你已经掌握了使用Python heapq模块构建优先队列和解决Top-K问题的核心方法。记住，默认的小顶堆通过heappop取出最小值，而通过维护一个大小为K的小顶堆是实现“前K大”问题的经典高效算法。