問題描述

怎樣從一個集合中獲得最大或最小的N個元素的列表？

解決方案

heapq模組有兩個函式：nlargest()和nsmallest()可以完美解決這個問題。

import heapq

nums = [1, 8, 2, 23, 7, -4, 18, 23, 42, 37, 2]
print(heapq.nsmallest(3, nums))
# 輸出結果：[-4, 1, 2]
print(heapq.nlargest(3, nums))
# 輸出結果：[42, 37, 23]

它們還可以接受一個關鍵字引數，用於更復雜的資料結構中：

import heapq

portfolio = [
    {'name': 'IBM', 'shares': 100, 'price': 91.1},
    {'name': 'AAPL', 'shares': 50, 'price': 543.22},
    {'name': 'FB', 'shares': 200, 'price': 21.09},
    {'name': 'HPQ', 'shares': 35, 'price': 31.75},
    {'name': 'YHOO', 'shares': 45, 'price': 16.35},
    {'name': 'ACME', 'shares': 75, 'price': 115.65}
]
# 以price的值進行比較
cheap = heapq.nsmallest(3, portfolio, lambda s: s['price'])
print(cheap)
"""
輸出結果：
[{'name': 'YHOO', 'shares': 45, 'price': 16.35}, {'name': 'FB', 'shares': 200, 'price': 21.09}, {'name': 'HPQ', 'shares': 35, 'price': 31.75}]
"""
 

討論

該用法的底層實現是堆排序

import heapq

nums = [1, 8, 2, 23, 7, -4, 18, 23, 42, 37, 2]
heapq.heapify(nums)
print(nums)
"""
輸出結果：
[-4, 2, 1, 23, 7, 2, 18, 23, 42, 37, 8]
"""

堆資料結構最重要的特徵是heap[0]永遠是最小的元素。並且剩餘的元素可以很容易的通過呼叫heapq.heappop()得到，該方法會先將第一個元素彈出，然後將下一個最小的元素排到第一個（這種操作的時間複雜度僅僅是O(logN)，N是堆大小）。比如，如果想要查詢最小的3個元素，可以這樣做：

import heapq

nums = [1, 8, 2, 23, 7, -4, 18, 23, 42, 37, 2]
heapq.heapify(nums)
print(heapq.heappop(nums), heapq.heappop(nums), heapq.heappop(nums))
"""
輸出結果：
-4 1 2
"""

總結

需要在不同的場合使用不同的方法，才能最大發揮它們的優勢。

• 當要查詢的元素個數N相對比較小的時候，函式nlargest()和nsmallest()是很合適的；
• 如果僅僅想要查詢唯一的（N=1）最小或最大元素，那麼使用函式min()和max()會更快些；
• 如果N的大小和集合大小接近，通常先排序再使用切片操作會更快些（sorted(items)[:N]或是sorted(items)[-N:]）。