问题

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

本文固定链接: http://www.js-code.com/python/python_50591.html