如何在 NumPy 数组中获取 N 个最大值的索引？

如何解决如何在 NumPy 数组中获取 N 个最大值的索引？

我能想到的最简单的是：

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提供了一种内置方法来进行部分排序；到目前为止，我还没有找到一个。

如果这个解决方案太慢（特别是对于 small n），可能值得考虑在Cython中编写一些东西。

解决方法

NumPy 提出了一种通过 获取数组最大值索引的方法np.argmax。

我想要类似的东西，但返回N最大值的索引。

例如，如果我有一个数组，[1,3,2,4,5],function(array,n=3)将返回[4,1]对应于元素的索引[5,3]。