问题
在某个集合中找到最大或最小的N个元素
解决方案
使用heapq
模块
heapq.nlargest(n, iterable, key=None)
heapq.nsmallest(n, iterable, key=None)
例如:
>>> import heapq
>>> l = [9, -2, 0, 8, 1, 3]
>>> print(heapq.nlargest(2, l))
[9, 8]
>>> print(heapq.nsmallest(2, l))
[-2, 0]
复制代码
此外,这两个函数都可以接受key
作为参数,例如:
import heapq
fruits = [
{'name': 'orange', 'price': 5},
{'name': 'apple', 'price': 2},
{'name': 'pear', 'price': 1.5},
{'name': 'lemon', 'price': 3},
]
print(heapq.nlargest(2, fruits, key=lambda x: x['price']))
复制代码
输出为:
[{'name': 'orange', 'price': 5}, {'name': 'lemon', 'price': 3}]
复制代码
讨论
根据Python3官方文档对heapq的介绍可以了解到
heapq
提供了堆数据结构的实现,并且实现方式是小顶堆,也就是说每次pop的时候取出的是最小的元素
首先使用heapq.heapify
将一个列表初始化为堆
>>> import heapq
>>> l = [-1, 2, 5, 0, 8]
>>> heapq.heapify(l)
>>> print(l)
[-1, 0, 5, 2, 8]
复制代码
然后就可以调用heapq.heappush
和heapq.heappop
对堆进行增加和删除操作了
>>> heapq.heappush(l, 8)
>>> print(l)
[-1, 0, 5, 2, 8, 8]
>>> print(heapq.heappop(l))
-1
复制代码
此外,heapq
还提供了其他堆的一些操作
heapq.heappushpop(heap, item)
先将item
存入堆中,然后弹出最小的元素,相当于先调用了heapq.heappush(item)
再调用heapq.heappop()
,但这样调用会比分开调用两个函数效率更高heapq.heapreplace(heap, item)
先弹出最小的元素,再存入item
heapq.merge(*iterables, key=None, reverse=False)
将多个有序的集合合并成一个有序的集合,并且返回的是迭代器对象
来源
Python Cookbook
关注
欢迎关注我的微信公众号:python每日一练