Python一天一模块: heapq 堆列队方法

heapq是一个内建模块, 实现了一个堆的数据结构，完美的解决了Top-K问题，以后解决Top-K问题的时候，直接把这个模块拿来用就可以了. 注意，默认的heap是一个小顶堆！

这个模块最有用的地方, 就是可以从列表等当中取出 最大的N个数或者最小的N个数.

• heapq.nlargest(n, iterable, key=None) 返回最大的n个元素（Top-K问题）
• heapq.nsmallest(n, iterable, key=None) 返回最小的n个元素（Top-K问题）
• heapq.heappush(heap, item) 把item添加到heap中（heap是一个列表）
• heapq.heappop(heap) 把堆顶元素弹出，返回的就是堆顶
• heapq.heappushpop(heap, item) 先把item加入到堆中，然后再pop，比heappush()再heappop()要快得多
• heapq.heapreplace(heap, item) 先pop，然后再把item加入到堆中，比heappop()再heappush()要快得多
• heapq.heapify(x) 将列表x进行堆调整，默认的是小顶堆
• heapq.merge(*iterables) 将多个列表合并，并进行堆调整，返回的是合并后的列表的迭代器

import heapq
import random
 
# Top-K
mylist = list(random.sample(range(100), 10))
k = 3
largest = heapq.nlargest(k, mylist)
smallest = heapq.nsmallest(k, mylist)
print('original list is', mylist)
print('largest-'+str(k), '  is ', largest)
print('smallest-'+str(k), ' is ', smallest)
 
# heapify
print('original list is', mylist)
heapq.heapify(mylist)
print('heapify  list is', mylist)
 
# heappush & heappop
heapq.heappush(mylist, 105)
print('pushed heap is', mylist)
heapq.heappop(mylist)
print('popped heap is', mylist)
 
# heappushpop & heapreplace
heapq.heappushpop(mylist, 130)    # heappush -> heappop
print('heappushpop', mylist)
heapq.heapreplace(mylist, 2)    # heappop -> heappush

在最常用的nlargest和nsmallest的方法当中, 支持key关键词, 即可自定义函数来实现排序处理:

portfolio = [
    {'name': 'IBM', 'shares': 100, 'price': 91.1},
    {'name': 'AAPL', 'shares': 50, 'price': 543.22},
    {'name': 'FB', 'shares': 200, 'price': 21.09},
    {'name': 'HPQ', 'shares': 35, 'price': 31.75},
    {'name': 'YHOO', 'shares': 45, 'price': 16.35},
    {'name': 'ACME', 'shares': 75, 'price': 115.65}
]
cheap = heapq.nsmallest(3, portfolio, key=lambda s: s['price'])
expensive = heapq.nlargest(3, portfolio, key=lambda s: s['price'])

heapq.heapify()方法可以排序数组, 首个和最末的元素就是最小和最大的. 可以根据需要取出.

但如果只要取最大或最小的数, 还是 min, max快点.

参考

• Py3官方文档: heap 列队算法