heapq模块

heapq模块中有两个函数--nlargest和nsmallest可以找到列表,元祖,集合等数据结构中的最大或最小的n个元素,并以列表的形式返回

import heapq

nums = {1, 8, 9, 41, 5, 6, -5, 21, 42, 37, 2}
# nums = [1, 8, 9, 41, 5, 6, -5, 21, 42, 37, 2]
# nums = (1, 8, 9, 41, 5, 6, -5, 21, 42, 37, 2)
print(heapq.nlargest(3, nums))
print(heapq.nsmallest(3, nums))

结果

[42, 41, 37]
[-5, 1, 2]

这两个函数都可以接受一个参数key,从而允许他们工作在更加复杂的数据结构之上

import heapq

portfolio = [
    {'name': 'Jan', 'age': 25, 'wages': 100},
    {'name': 'ming', 'age': 10, 'wages': 55},
    {'name': 'san', 'age': 32, 'wages': 321.5},
    {'name': 'si', 'age': 28, 'wages': 222.2},
    {'name': 'wu', 'age': 43, 'wages': 256},
]
cheap = heapq.nsmallest(3, portfolio, key=lambda s: s['wages'])
expensive = heapq.nlargest(3, portfolio, key=lambda s: s['wages'])
print(cheap)
print(expensive)

结果

[{'age': 10, 'name': 'ming', 'wages': 55}, {'age': 25, 'name': 'Jan', 'wages': 100}, {'age': 28, 'name': 'si', 'wages': 222.2}]
[{'age': 32, 'name': 'san', 'wages': 321.5}, {'age': 43, 'name': 'wu', 'wages': 256}, {'age': 28, 'name': 'si', 'wages': 222.2}]

如果寻找最大或最小的n个元素,且同数据结构中的元素相比,n很小,那么上述的方法性能不是那么的乐观,这样的情况下我们可以使用堆

import heapq

nums = [1, 4, 5, 7, 36, -3, 65, 33, 26, 28, 2, 4]
heap = list(nums)
heapq.heapify(heap)
print(heap)

结果

[-3, 2, 1, 7, 4, 4, 65, 33, 26, 28, 36, 5]

这些函数首先会在底层将数据转化成列表,且元素会以堆的形式排列

堆最重要的性能就是heap[0]总是最小的那个元素,接下来的元素可依次通过heapq.heappop()方法找到

该方法会将第一个元素(最小的)弹出,然后以第二个元素取而代之(这个操作的复杂度是O(logN),N代表堆的大小)