python中的优先队列

在 Python 中，优先队列（Priority Queue）是一个可以随时获取队列中最大（或最小）元素的数据结构。Python 的标准库 heapq 提供了一个实现最小堆的优先队列，默认情况下是最小堆，但可以通过一些技巧来实现最大堆。优先队列在算法中常用于求解最短路径、合并有序链表、求解 k 个最小/最大的元素等问题。

heapq 模块的基本用法

heapq 模块提供了一些函数来操作堆（即优先队列），堆在 Python 中实际上是基于列表实现的。

1. heapq.heappush(heap, item)

• 将元素 item 加入到堆 heap 中，保持堆的性质。

2. heapq.heappop(heap)

• 弹出并返回堆中的最小元素，同时保持堆的性质。

3. heapq.heapify(iterable)

• 将一个可迭代对象 iterable 转换成堆，原地操作。

4. heapq.heappushpop(heap, item)

• 将元素 item 加入堆中，并弹出并返回堆中的最小元素（两步操作合并为一步，性能较高）。

5. heapq.heapreplace(heap, item)

• 弹出堆中的最小元素，并将新元素 item 插入堆中，保持堆的性质。与 heappop 和 heappush 连续操作相比，它性能更高。

6. heapq.nlargest(n, iterable)

• 返回 iterable 中的前 n 个最大元素，时间复杂度为 (O(n \log k))。

7. heapq.nsmallest(n, iterable)

• 返回 iterable 中的前 n 个最小元素，时间复杂度为 (O(n \log k))。

---

最大堆的实现

Python 的 heapq 模块默认是最小堆。如果要实现最大堆，可以通过将所有元素取反来模拟。如下所示：

import heapq

# 模拟最大堆
max_heap = []
heapq.heappush(max_heap, -10)
heapq.heappush(max_heap, -5)
heapq.heappush(max_heap, -20)

# 弹出最大值
max_val = -heapq.heappop(max_heap)  # 输出 20
print(max_val)

通过将所有元素存储为负值，堆的最小元素即为我们想要的最大元素，从而实现了最大堆的功能。

---

优先队列的应用场景

1. 求解最小值或最大值问题
   通过优先队列可以高效地从动态数据中提取出最小值或最大值。

2. Dijkstra 最短路径算法
   优先队列常用于 Dijkstra 算法中，从中选择当前最小的路径。

3. 合并多个已排序的数组或链表
   用优先队列可以高效地将多个有序链表合并为一个有序链表。

4. Top K 问题
   在求解给定数组的前 K 大或前 K 小元素时，优先队列是一个非常有效的工具。

---

示例代码：合并多个已排序的链表

这是一个经典的优先队列应用题目。我们需要合并多个已排序的链表，可以使用优先队列来逐个提取链表头部的最小值。

import heapq

# Definition for singly-linked list.
class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next

    def __lt__(self, other):
        return self.val < other.val  # 用于堆排序时对 ListNode 排序
    
def mergeKLists(lists):
    min_heap = []
    # 将所有链表的头节点加入堆
    for l in lists:
        if l:
            heapq.heappush(min_heap, l)
    
    # 哨兵节点，用来拼接结果链表
    head = point = ListNode(0)
    
    while min_heap:
        # 取出当前堆中最小的元素
        node = heapq.heappop(min_heap)
        point.next = node
        point = point.next
        if node.next:
            heapq.heappush(min_heap, node.next)  # 将下一个节点加入堆

    return head.next

在这个例子中，我们定义了 ListNode 类，并使用 heapq 实现了合并多个链表的操作。由于 ListNode 对象没有默认的排序规则，我们需要为其实现 __lt__ 方法，用来支持堆排序。

---

总结

在算法题中，优先队列的应用非常广泛，特别是需要动态获取最大/最小元素时。heapq 提供了一个最小堆实现，通过巧妙地使用可以实现最大堆或其他类型的优先队列。掌握优先队列的基本操作和应用，可以帮助你高效地解决很多类型的问题。