python中heapq的库是什么_python中heapq模块详解和简单应用

问题：怎样从一个集合中获得最大或者最小的 N 个元素列表？

含义：从一个列表或者集合中获取N个最小元素，就等价于将一个列表排好序，提取出最小的一部分或者最大的 一部分。下面就举例子模拟一下原理。

1grades =[98,90,89,100,87,96,94,100]

2grades.sort() #在原来基础上排序

3print(grades) #打印[87, 89, 90, 94, 96, 98, 100, 100]

4#查找最大的三个数就等同于将排序好的列表提取最大的三个数

5#很明显这里最大的三个数是[98,100,100]

6print(grades[-3:])

上面是模拟了下问题的含义，这不是我们的本意，在这里我们将采用其他的方法提取最大或者最小的 N 个元素列表。

标准库有个heapq模块，heapq 模块有两个函数：nlargest() 和 nsmallest() 可以完美解决这个问题。

heapq，英文单词含义：堆排序算法，heap是堆的意思，在这里是数据堆结构。

1#导入堆排序模块

2import heapq

3grades =[87, 89, 90, 94, 96, 98, 100, 100]

4import random

5random.shuffle(grades) #打乱正序

6#print(grades)

7three_max = heapq.nlargest(3,grades)

8three_min = heapq.nsmallest(3,grades)

9print(f'最大的三个数{three_max},最小的三个数{three_min}')

上面的random.shuffle用于打乱列表顺序，这里可有可无，这里就是为了使样本列表无序。heapq.nlargest还有第三个参数就是key，类似于sorted(iterable,key=function)，这里同样有关键字key情况下，把key绑定函数，下面用lambda函数作为key的排序依据参数：

1#有关键字key情况下

2portfolio = [

3{'name': 'IBM', 'shares': 100, 'price': 91.1},

4{'name': 'AAPL', 'shares': 50, 'price': 543.22},

5{'name': 'FB', 'shares': 200, 'price': 21.09},

6{'name': 'HPQ', 'shares': 35, 'price': 31.75},

7{'name': 'YHOO', 'shares': 45, 'price': 16.35},

8{'name': 'ACME', 'shares': 75, 'price': 115.65}

9]

10#找出价格便宜的3个，和价格最贵的三个

11cheap = heapq.nsmallest(3,portfolio,key=lambda x : x["price"]) #返回的是列表形式

12expensive = heapq.nlargest(3,portfolio,key=lambda y : y["price"])

13print(f"最便宜的三家:{cheap}\n最贵的三家:{expensive}")

上面将字典按照price值的大小进行排序，分别取出3个最大值和3个最小值。

如果你想在一个集合中查找最小或最大的 N 个元素，并且 N 小于集合元素数量，那么这些函数提供了很好的性能。因为在底层实现里面，首先会先将集合数据进行堆排序(找出最小值放在最底下，其他元素基本不变)后放入一个列表中。

堆数据结构最重要的特征是 heap[0] 永远是最小的元素。并且剩余的元素可以很容易的通过调heapq.heappop() 方法得到，该方法会先将第一个元素弹出来，然后用下一个最小的元素来取代被弹出元素。

1import heapq

2data =  [1, 8, 2, 23, 7, -4, 18, 23, 42, 37, 2]

3heap = list(data)

4heapq.heapify(heap) #进行堆排序，data[0]永远都是最小的一个，其他元素不管，data成为堆数据结构

5print(heap) #[-4, 2, 1, 23, 7, 2, 18, 23, 42, 37, 8]

6v = heapq.heappop(heap)

7print(v,heap)

8v = heapq.heappop(heap) #从左侧移除一个元素后，剩下元素自动变成最小的在最左侧。

9print(v,heap)

10print([heapq.heappop(heap) for i in range(3)])

11heapq.heappush(heap,-100) #再往堆结构里添加一个数据

12print(heapq.heappop(heap)) #print -100

这种方法适用于从堆排序取出最小数值，且不需要进行整个列表的排序。

当要查找的元素个数相对比较小的时候，函数 nlargest() 和 nsmallest() 是很合适的。如果仅仅想查找唯一的最小或最大的元素，那么使用 min() 和max() 函数会更快些。类似的，如果查询元素个数N 的大小和集合大小接近的时候，通常先排序这个集合然后再使用切片操作会更快点( sorted(items)[:N] 或者是 sorted(items)[-N:])。需要在正确场合使用函数nlargest() 和 nsmallest() 才能发挥它们的优势(如果N 快接近集合大小了，那么使用排序操作会更好些)。