Python每日一練0006

問題

在某個集合中找到最大或最小的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 heapqfruits = [ {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每日一練

http://weixin.qq.com/r/YygOFnnE8iSNrT11931x (二維碼自動識別)

推薦閱讀：