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 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每日一練