Python中使用heapq模块找到最大或最小的N个元素

Python中使用heapq模块找到最大或最小的N个元素

在heapq中有两个函数模块，分别是nlargest()、nsmallest(),可以用这两个函数分别找到列表中最大和最小的元素。

这两个函数分别有三个参数，分别为n, iterable, key

1. n: 限制需要的元素个数
2. iterable：可迭代对象
3. key：对复杂的数据进行处理

情况一：仅对列表数据进行操作

对列表进行处理时无需考虑key值，仅需限定参数n、iterable即可。

'''
Author：爱吃跳跳糖的小怪兽
time：2020-8-27
Version：1.0
'''

import heapq

List_Example1 = [-100, 2, 5, 7, -78, 85, 56, -21, 0]     #假定一个无序的整数列表需要进行数据提取
print(heapq.nlargest(3,List_Example1))
print(heapq.nsmallest(3,List_Example1))

Pycharm中运行结果如下：

情况二：需要对稍微复杂的数据进行操作

如果不设定key值时会出现以下错误

代码块

'''
Author：爱吃跳跳糖的小怪兽
time：2020-8-27
Version：1.0
'''

import heapq
achievements = [
	{'name':'xu', 'data':98},
	{'name':'yuan', 'data':95},
	{'name':'qi', 'data':99},
	{'name':'shuai', 'data':45},
	{'name':'ge', 'data':70}
]
print(heapq.nsmallest(3,achievements))
print(heapq.nlargest(3,achievements))

报错截图：

错误分析：程序抛出了一个类型错误，表示不支持对两个字典进行比较。

解决方法：对key值进行设定，这里使用lamda匿名函数将data键对应的值传返回至key值中，比对各个字典中的data数据以进行排列。

'''
Author：爱吃跳跳糖的小怪兽
time：2020-8-27
Version：1.0
'''

import heapq

achievements = [
	{'name':'xu', 'data':98},
	{'name':'yuan', 'data':95},
	{'name':'qi', 'data':99},
	{'name':'shuai', 'data':45},
	{'name':'ge', 'data':70}
]
print(heapq.nsmallest(3,achievements,key=lambda x:x['data']))
print(heapq.nlargest(3,achievements,key=lambda x:x['data']))

输出结果截图：

总结

在使用heapq模块的nlargest()和nsmallest()函数进行数据查找时，需要先行判断数据类型，考虑是否设定key值，以确保函数的正常运行。