首页 > 代码库 > 【Python】Python取top N相关的模块：heapq模块

【Python】Python取top N相关的模块：heapq模块

2024-07-05 09:30:58   222人阅读

最近在程序中需要取一个列表的top 3元素，就是去一个列表中数值最大的3个元素。这可以用Python的heapq模块来处理。

1、对列表取top N：

现在有一个列表myList，需要取出该列表的最大3个元素和最小3个元素，按如下代码所述的简单例子：
test.py
import heapq

myList = [5, 2, 6, 12, 7, 3, 4, 9]
topNum = 3 
nlargestList = heapq.nlargest(topNum, myList)       #取最大3个元素
nsmallestList = heapq.nsmallest(topNum, myList)     #取最小3个元素
print nlargestList                                  #输出结果
print nsmallestList
我是在Linux环境下，输入：python test.py，就可以运行程序了。

输出结果：
[12, 9, 7]
[2, 3, 4]

没错，这也可以用排序来获得，但是对另外一些情况，排序就比较麻烦了，比如对字典取top N的时候，因为其本身就是无序的，即使排序也不会改变元素在字典中的位置。

2、对字典取top N：

2.1 按字典的key值取top N：

现在有一个字典，想按字典的key值取key值最大的3个元素和key值最小的3个元素，如下代码是一个简单的例子：
test.py
import heapq

myDict = {1:'a', 3:'d', 8:'g', 5:'m', 9:'t', 4:'s', 2:'u'}   #定义一个字典
topNum = 3
nlargestList = heapq.nlargest(topNum, myDict.keys())         #取key值最大的3个元素
nsmallestList = heapq.nsmallest(topNum, myDict.keys())       #取key值最小的3个元素
for key in nlargestList:                                     #输出结果
    print key,myDict[key]
for key in nsmallestList:
    print key,myDict[key]
我是在Linux环境下，输入：python test.py，就可以运行程序了。

输出结果：
9 t
8 g
5 m
1 a
2 u
3 d

从结果可以看出，确实符合要求。

2.2 按字典的value值取top N：

现在有一个字典，想按字典的value值取value值最大的3个元素和value值最小的3个元素，如下一个简单的例子：
test.py
import heapq

myDict = {1:'a', 3:'d', 8:'g', 5:'m', 9:'t', 4:'s', 2:'u'}
topNum = 3
nlargestList = heapq.nlargest(topNum, myDict.values())    #取value值最大的3个元素，注意与上一程序区别
nsmallestList = heapq.nsmallest(topNum, myDict.values())  #取value值最小的3个元素，注意与上一程序区别
for value in nlargestList:                                #输出结果
    for key in myDict:
        if myDict[key] == value:
            print key,myDict[key]
for value in nsmallestList:                               #输出结果
    for key in myDict:
        if myDict[key] == value:
            print key,myDict[key]
我是在Linux环境下，输入：python test.py，就可以运行程序了。

输出结果：
2 u
9 t
4 s
1 a
3 d
8 g

从结果可以看出，确实符合要求。

希望对大家有所帮助。


联系
我们