问题

在某个集合中找到最大或最小的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每日一练