NumPy 提出了一种通过 获取数组最大值索引的方法np.argmax。
np.argmax
我想要类似的东西,但返回N最大值的索引。
N
例如,如果我有一个数组,[1, 3, 2, 4, 5],function(array, n=3)将返回[4, 3, 1]对应于元素的索引[5, 4, 3]。
[1, 3, 2, 4, 5]
function(array, n=3)
[4, 3, 1]
[5, 4, 3]
我能想到的最简单的是:
In [1]: import numpy as np In [2]: arr = np.array([1, 3, 2, 4, 5]) In [3]: arr.argsort()[-3:][::-1] Out[3]: array([4, 3, 1])
这涉及到一个完整的数组排序。我想知道是否numpy提供了一种内置方法来进行部分排序;到目前为止,我还没有找到一个。
numpy
如果这个解决方案太慢(特别是对于 small n),可能值得考虑在Cython中编写一些东西。
n