python list中找Topk的数值和索引

转载自https://blog.csdn.net/lwgkzl/article/details/107484265

需求：
对于一个python list 或者numpy数组，我需要找到这个list中最大的K个数及其对应的下标。

解决方式：

1. 可以构造字典通过排序解决，不过代码量较多。

2. 使用heapq库，可以直接获取最大值的下标和数值。

import heapq
a = [4,2,6,1,9,9]
 
# 获取下标， 输出为[4, 5, 2]
heapq.nlargest(3, range(len(a)), a.__getitem__)
 
# 获取数值， 输出为[9, 9, 6]
heapq.nlargest(3,a)

如果要取最小的数，使用 nsmallest即可。