0.摘要

本文主要介紹heapq模組提供的主要函式方法。

1.nlargest()和nsmallest

從字面意思就可以知道

heapq.nlargest(n, iterable, key=None) ：返回可列舉物件中的 n 個最大值
heapq.nsmallest(n, iterable, key=None) ：返回可列舉物件中的 n 個最小值。

import heapq

nums = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
print(heapq.nlargest(3,nums))
print(heapq.nsmallest(3,nums))
#result:[9, 8, 7]
#result:[0, 1, 2]
 

相比於max()和min()函式，nlargest()和nsmallest()能夠按遞增或遞減順序給出最大或最小的n個值。

另外，nlargest()和nsmallest()更強大的地方在於能夠接受一個引數key，從而允許了這兩個函式能夠工作在更為複雜的資料結構上。

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}
]

cheap = heapq.nsmallest(3, portfolio, key=lambda s: s['price'])
expensive = heapq.nlargest(3, portfolio, key=lambda s: s['price'])

print(cheap)
print(expensive)
#result:[{'name': 'YHOO', 'shares': 45, 'price': 16.35}, {'name': 'FB', 'shares': 200, 'price': 21.09}, {'name': 'HPQ', 'shares': 35, 'price': 31.75}]
#result:[{'name': 'AAPL', 'shares': 50, 'price': 543.22}, {'name': 'ACME', 'shares': 75, 'price': 115.65}, {'name': 'IBM', 'shares': 100, 'price': 91.1}]
 

通過key引數，可以以字典中的某個量為依據，找出最大或最小的n個值。

注意：key引數並非一定是數值，字串也同樣可以（按照ASCII碼排序）。

提示：max()和min()函式也具有key引數，這裡替換為max()和min()函式即可得到最大和最小的值。

2.heappush()和heappop

heapq.heappush(heap, item)：將元素新增到heap中
heapq.heappop(heap)：返回 root 節點，即 heap 中最小的元素

使用heappush()和heappop實現一個優先佇列例子，即每次pop操作，彈出的是優先順序最大的與元素。

注意：heapq.heappop(）彈出的是優先順序最小的元素，這裡通過對優先順序取相反數顛倒優先順序。

import heapq

class PriorityQueue:
    def __init__(self):
        self._queue = []
        self._index = 0

    def push(self, item, priority):
        heapq.heappush(self._queue, (-priority, self._index, item))
        self._index += 1

    def pop(self):
        return heapq.heappop(self._queue)[-1]

# Example use
class Item:
    def __init__(self, name):
        self.name = name
    def __repr__(self):
        return 'Item({!r})'.format(self.name)

q = PriorityQueue()
q.push(Item('foo'), 1)
q.push(Item('bar'), 5)
q.push(Item('spam'), 4)
q.push(Item('grok'), 1)

print("Should be bar:", q.pop())
print("Should be spam:", q.pop())
print("Should be foo:", q.pop())
print("Should be grok:", q.pop())
#result:Should be bar: Item('bar')
#result:Should be spam: Item('spam')
#result:Should be foo: Item('foo')
#result:Should be grok: Item('grok')

這裡引入index是為了防止優先順序相同的情況，根據index進行排序。

元組(-priority, item)在priority相同的情況下，元組會對item進行比較。

但是，Item例項在python中是不能進行比較運算的，強制比較會導致報錯。

加入index就可以避免這一問題，因為每一個元素的index都是不同的。這樣，在priority相同的情況下，會比較index大小，並且由於index必然不相等，所以不存在進一步比較item的情況，從而避免了報錯。