python heapq模块

概述

python有一个内置的heapq模块，对应着数据结构——堆。堆是一个很重要的数据结构，可以看做一棵完全二叉树的数组对象。

一个数组，逻辑上是可以根据索引构造出一棵完全二叉树的，其节点位置和索引的关系为（索引从0开始）：

• $i$位置结点的左孩子在数组中的位置是：$(2\ast i)+1$，右孩子在数组中的位置是：$(2\ast i)+2$；
• $i$位置结点的父节点在数组中的位置是：$(i-1)//2$；

这也就是堆特征。堆的类型有两种：大根堆和小根堆。顾名思义，大根堆就是父节点大于子节点，小根堆就是父节点小于子节点，python内置的heapq模块默认的是小根堆。

模块重要函数

heapq模块的一些重要函数：

• heapify(heap)
• heappush(heap, x)
• heappop(heap)
• heapreplace(heap, x)
• nlargest(n, iter) 和 nsmallest(n, iter)

heapify(heap)

该函数的作用是让列表具备堆特征：

import heapq

arr = [8, 5, 3, 9, 1, 4, 2, 0, 7, 6]
heapq.heapify(arr)
print(arr)
# 此时arr为：[0, 1, 2, 5, 6, 4, 3, 9, 7, 8]

在构建堆时，从索引为$n/2-1$处开始往上构建，以上面的例子来说，便是从数字“1”开始。

heappush(heap, x)

heappush就是将数字x添加到堆中：

import heapq

arr = [8, 5, 3, 9, 1, 4, 2, 0, 7, 6]
heapq.heapify(arr)
print(arr)
# 此时arr为：[0, 1, 2, 5, 6, 4, 3, 9, 7, 8]
heapq.heappush(arr, 4.5)
print(arr)
# 此时arr为：[0, 1, 2, 5, 4.5, 4, 3, 9, 7, 8, 6]

heappop(heap)

heappop函数便是弹出堆顶元素，并将剩下的堆继续保持堆特征即最小的数依旧在索引0处：

import heapq

arr = [8, 5, 3, 9, 1, 4, 2, 0, 7, 6]
heapq.heapify(arr)
print(arr)
# 此时arr为：[0, 1, 2, 5, 6, 4, 3, 9, 7, 8]
heapq.heappush(arr, 4.5)
print(arr)
# 此时arr为：[0, 1, 2, 5, 4.5, 4, 3, 9, 7, 8, 6]
heapq.heappop(arr)
print(arr)
# 此时arr为：[1, 4.5, 2, 5, 6, 4, 3, 9, 7, 8]

heapreplace(heap, x)

heapreplace的用处是从堆中弹出最小的元素，再压入一个新元素。相比于依次执行函数heappop和heappush：

import heapq

arr = [8, 5, 3, 9, 1, 4, 2, 0, 7, 6]
heapq.heapify(arr)
print(arr)
# 此时arr为：[0, 1, 2, 5, 6, 4, 3, 9, 7, 8]
heapq.heappush(arr, 4.5)
print(arr)
# 此时arr为：[0, 1, 2, 5, 4.5, 4, 3, 9, 7, 8, 6]
heapq.heappop(arr)
print(arr)
# 此时arr为：[1, 4.5, 2, 5, 6, 4, 3, 9, 7, 8]
heapq.heapreplace(arr, 5.5)
print(arr)
# 此时arr为：[2, 4.5, 3, 5, 6, 4, 5.5, 9, 7, 8]

模块heap中还有一个函数为heappushpop(heap, x)，此函数便相当于先执行以此heappush，再执行一次heappop操作，与heapreplace类似。

nlargest(n, iter) 和 nsmallest(n, iter)

顾名思义，这两个函数便是输出最大$n$个值和输出最小$n$个值：

import heapq

arr = [8, 5, 3, 9, 1, 4, 2, 0, 7, 6]
heapq.heapify(arr)
print(arr)
# 此时arr为：[0, 1, 2, 5, 6, 4, 3, 9, 7, 8]
heapq.heappush(arr, 4.5)
print(arr)
# 此时arr为：[0, 1, 2, 5, 4.5, 4, 3, 9, 7, 8, 6]
heapq.heappop(arr)
print(arr)
# 此时arr为：[1, 4.5, 2, 5, 6, 4, 3, 9, 7, 8]
heapq.heapreplace(arr, 5.5)
print(arr)
# 此时arr为：[2, 4.5, 3, 5, 6, 4, 5.5, 9, 7, 8]
print(heapq.nlargest(4, arr))
# 输出为：[9, 8, 7, 6]
print(heapq.nsmallest(4, arr))
# 输出为：[2, 3, 4, 4.5]

此处的iter参数不仅可以输入具有堆特征的数据（比如例中heapify后的arr），也可以就输入一个列表或其他任意可以迭代的数据结构：

import heapq

items = [
    {'item': 'a', 'age': 23},
    {'item': 'b', 'age': 11},
    {'item': 'c', 'age': 66},
    {'item': 'd', 'age': 15},
    {'item': 'e', 'age': 45},
    {'item': 'f', 'age': 52}
]

top3 = heapq.nlargest(3, items, key=lambda x:x['age'])
print(top3)
# 输出为：[{'item': 'c', 'age': 66}, {'item': 'f', 'age': 52}, {'item': 'e', 'age': 45}]