Python heapq模块：最小堆的使用

堆是非线性的树形的数据结构（完全二叉树），有两种堆，最大堆与最小堆。（ heapq库中的堆默认是最小堆）

• 最大堆，树中各个父节点的值总是大于或等于任何一个子节点的值。
• 最小堆，树中各个父节点的值总是小于或等于任何一个子节点的值。

我们一般使用二叉堆来实现优先级队列，它的内部调整算法复杂度为$\log N$。堆是一个二叉树，其中最小堆每个父节点的值都小于或等于其所有子节点的值。整个最小堆的最小元素总是位于二叉树的根节点。

Python的heapq模块提供了对最小堆的建立和使用。heapq堆数据结构最重要的特征是heap[0]永远是最小的元素。

heapq模块中的常用方法：

下面来一一讲解这些方法：

使用heapq创建堆

heapq中创建堆的方法有两种。

• heappush(heap, num)，先创建一个空堆，然后将数据一个一个地添加到堆中。每添加一个数据后，heap都满足小顶堆的特性。
• heapify(array)，直接将数据列表调整成一个小顶堆。

两种方法实现的结果会有差异，但都满足小顶堆的特性，不影响堆的使用（堆只会从堆顶开始取数据，取出数据后会重新调整结构）。

heapq.heappush(heap, item)

heap为定义堆，item增加的元素

>>> h = []
>>> heapq.heappush(h, 2)
>>> h
[2]
1
2
3
4

heapq.heapify(list)

将列表转换为堆

>>> num = [1, 2, 34, 5, 6, 67, 3, 2, 5, 6, 4]
>>> heapq.heapify(num)
>>> num
[1, 2, 3, 2, 4, 67, 34, 5, 5, 6, 6]
1
2
3
4

堆顶数据出堆的方法

heapq.heappop(heap)

将堆顶的数据出堆，并将堆中剩余的数据构造成新的小顶堆。也即是：删除并返回最小值，因为堆的特征是heap[0]永远是最小的元素，所以一般都是删除第一个元素。

>>> num
[1, 2, 3, 2, 4, 67, 34, 5, 5, 6, 6]
>>> heapq.heappop(num)
1
>>> num
[2, 2, 3, 5, 4, 67, 34, 5, 6, 6]
1
2
3
4
5
6

heapq替换数据的方法

heapq.heapreplace(heap, item)

删除并返回最小元素值，添加新的元素值

>>> num = [1, 2, 3, 4, 5, 6, 5, 8, 9]
>>> heapq.heapreplace(num, 99)
1
>>> num
[2, 4, 3, 8, 5, 6, 5, 99, 9]
1
2
3
4
5

heapq.heappushpop(list, item)

判断添加元素值与堆的第一个元素值对比：

• 如果大，则删除并返回第一个元素，然后添加新元素值item
• 如果小，则返回item，原堆不变。

>>> num
[2, 4, 3, 8, 5, 6, 5, 99, 9]
>>> heapq.heappushpop(num, 6)
2
>>> num
[3, 4, 5, 8, 5, 6, 6, 99, 9]
>>> heapq.heappushpop(num, 1)
1
>>> num
[3, 4, 5, 8, 5, 6, 6, 99, 9]
1
2
3
4
5
6
7
8
9
10

总结：

• heappushpop(heap, num)，先将num添加到堆中，然后将堆顶的数据出堆。
• heapreplace(heap, num)，先将堆顶的数据出堆，然后将num添加到堆中。

两个方法都是即入堆又出堆，只是顺序不一样，可以用于替换堆中的数据。

使用heapq合并两个有序列表

heapq.merge(list1, list2)

将两个有序的列表合并成一个新的有序列表，返回结果是一个迭代器。这个方法可以用于归并排序。

>>> num
[3, 4, 5, 8, 5, 6, 6, 99, 9]
>>> num2
[1000]
>>> list(heapq.merge(num2, num))
[3, 4, 5, 8, 5, 6, 6, 99, 9, 1000]
1
2
3
4
5
6

应用：合并两个有序列表

>>> a = [10, 7, 15, 8]
>>> b = [17, 3, 8, 20, 13]
>>> merge = heapq.merge(sorted(a), sorted(b))
>>> list(merge)
[3, 7, 8, 8, 10, 13, 15, 17, 20
1
2
3
4
5

获取堆中的最小值或最大值

• nlargest(num, heap)，从堆中取出num个数据，从最大的数据开始取，返回结果是一个列表（即使只取一个数据）。如果num大于等于堆中的数据数量，则从大到小取出堆中的所有数据，不会报错，相当于实现了降序排序。
• nsmallest(num, heap)，从堆中取出num个数据，从最小的数据开始取，返回结果是一个列表。

这两个方法除了可以用于堆，也可以直接用于列表，功能一样。

heapq.nlargest(n,heap)

查询堆中的最大的n个元素

>>> num
[3, 4, 5, 8, 5, 6, 6, 99, 9]
>>> heapq.nlargest(3, num)
[99, 9, 8]
1
2
3
4

heapq.nsmallest(n,heap)

查询堆中的最小n个元素

>>> num
[3, 4, 5, 8, 5, 6, 6, 99, 9]
>>> heapq.nsmallest(3, num)
[3, 4, 5]
1
2
3
4

使用heapq实现堆排序

先将待排序列表中的数据添加到堆中，构造一个小顶堆。然后依次将堆顶的值取出，添加到一个新的列表中，直到堆中的数据取完，新列表就是排序后的列表。

heappop(heap)，将堆顶的数据出堆，并将堆中剩余的数据构造成新的小顶堆。

>>> array = [10, 17, 50, 7, 30, 24, 27, 45, 15, 5, 36, 21]
>>> heap = heapq.heapify(array)
>>> [heapq.heappop(array) for i in range(len(array))]
[5, 7, 10, 15, 17, 21, 24, 27, 30, 36, 45, 50]
1
2
3
4

优先级队列

class PriorityQueue:
    def __init__(self):
        self._queue = []
        self._index = 0
        
    def push(self, item, priority):
        heapq.heappush(self._queue, (-priority, self._index, item))
        # 第一个参数：添加进的目标序列
        # 第二个参数：将一个元组作为整体添加进序列，目的是为了方便比较
        # 在priority相等的情况下，比较_index
        # priority为负数使得添加时按照优先级从大到小排序，因为堆排序的序列的第一个元素永远最小
        self._index += 1
        
    def pop(self):
        # 返回按照-priority 和_index排序后的第一个元素(是一个元组)的最后一个元素(item)
        return heapq.heappop(self._queue)[-1]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

>>> q = PriorityQueue()
>>> q.push("bar", 2)
>>> q.push("foo", 1)
>>> q.push("gork", 3)
>>> q.push("new", 1)
>>> print(q.pop())
gork  # 优先级最高
>>> print(q.pop())
bar   # 优先级第二
>>> print(q.pop())
foo   # 优先级与new相同，比较index，因为先加入，index比new小，所以排在前面
>>> print(q.pop())
new 
1
2
3
4
5
6
7
8
9
10
11
12
13

\quad
\quad

THE END