博客

Python中的优先队列(priority queue)和堆(heap)

  • Post category:Python

1. 什么是优先队列和堆

优先队列 (Priority Queue) 是一种特殊的队列，元素带有优先级，高优先级的元素可以先被取出来。常见的实现方式包括堆和二叉搜索树等。

堆 (Heap) 是一种完全二叉树，它分为两类，分别是小根堆和大根堆，小根堆中，每个节点的值不大于其父节点的值，大根堆中，每个节点的值不小于其父节点的值。堆常见的操作有插入一个元素和弹出堆顶元素两种。

2. Python中的优先队列和堆

Python中的heapq模块实现了堆的操作，能够很方便地对列表等数据结构进行堆化操作，支持普通的堆和优先队列。下面我们就来看一下Python中如何实现优先队列。

查看Python文档中heapq模块的介绍能够发现该模块支持heap的基本操作，但不是真正意义上的priority queue。那我们该如何实现priority queue呢？其实，我们可以通过定义一个二元组 (priority_number，item) 来实现。其中，priority_number 是用来比较优先级的关键值，在元组中排在 item 前面。这样，我们就可以根据 priority_number 的值来进行排序，实现 priority queue。

以下是一个简单的示例代码：

import heapq

class PriorityQueue:
    def __init__(self):
        self._queue = []
        self._index = 0
​
    def push(self, item, priority):
        heapq.heappush(self._queue, (-priority, self._index, item))
        self._index += 1

    def pop(self):
        return heapq.heappop(self._queue)[-1]

在这个示例中，我们定义了一个 PriorityQueue 类。其中，push 方法和 pop 方法就是我们实现的优先队列操作。我们调用了 Python 的 heapq 模块中的 heappush 和 heappop 方法，将 (priority_number，_index，item) 这个元组进行入队和出队操作。因为 PriorityQueue 类实现了自己的 _index 属性，我们可以在元组中添加 _index 来确保队列按照入队顺序排序。通过将 priority_number 取反，可以实现优先级递增的效果。

我们来看一个例子，创建队列并添加元素：

q = PriorityQueue()
q.push("hello", 5)
q.push("world", 1)
q.push("python", 4)
q.push("code", 3)

这个队列的定义方式和其它队列一样，通过 push 方法添加元素。每个元素包括一个字符串和一个数字，数字代表这个元素的优先级。我们可以打印出这个队列的内容：

while q._queue:
    print(q.pop())

这样就可以看到输出结果是按照优先级降序排列的了。

3. Python中使用heap实现堆

接下来我们看一下如何使用 Python 的 heapq 模块实现堆。以下是一个简单的示例代码：

import heapq

# 生成一个heap，并将数组转换为小根堆
a = [3, 4, 2, 1, 5, 6, 7]
heapq.heapify(a)

# 弹出堆顶元素，并返回弹出的值
print(heapq.heappop(a))

# 向堆中插入元素
heapq.heappush(a, 0)

# 取出堆中最小的 n 个元素(不弹出)
print(heapq.nsmallest(2, a))
# 取出堆中最大的 n 个元素(不弹出)
print(heapq.nlargest(2, a))

在这个示例中，我们使用 Python 的 heapq 模块生成了一个小根堆。首先我们弹出了堆顶元素，然后向堆中插入了一个元素。最后，我们取出堆中最小的 2 个元素和最大的 2 个元素。我们可以看到在 Python 中使用 heapq 模块实现堆是非常简单的。

总结：

本文介绍了Python中的优先队列和堆。首先我们讲解了优先队列和堆的定义，并且说明了优先队列可以通过Python的heap模块实现。接下来我们结合示例代码介绍了如何实现自己的priority queue，以及如何使用heap实现堆。使用这些方法，我们可以非常方便地对数据进行排序和处理。

Tags: