對於一個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即可