对于一个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即可