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

Question

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

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

For instance, if I have an array, [1, 3, 2, 4, 5] , then nargmax(array, n=3) would return the indices [4, 3, 1] which correspond to the elements [5, 4, 3] .

原文由 Alexis Métaireau 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 1

较新的 NumPy 版本（1.8 及更高版本）为此有一个名为 argpartition 的函数。要获得四个最大元素的索引，请执行

>>> a = np.array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])
>>> a
array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])

>>> ind = np.argpartition(a, -4)[-4:]
>>> ind
array([1, 5, 8, 0])

>>> top4 = a[ind]
>>> top4
array([4, 9, 6, 9])

与 argsort 不同，此函数在最坏的情况下以线性时间运行，但返回的索引未排序，从评估结果可以看出 a[ind] 。如果您也需要，请在之后对它们进行排序：

 >>> ind[np.argsort(a[ind])]
array([1, 8, 5, 0])

以这种方式按排序顺序获取前 k 个元素需要 O( n + k log k ) 时间。

原文由 Fred Foo 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 2

我能想到的最简单的是：

 >>> import numpy as np
>>> arr = np.array([1, 3, 2, 4, 5])
>>> arr.argsort()[-3:][::-1]
array([4, 3, 1])

这就涉及到对数组的完整排序。我想知道 numpy 是否提供了一种内置的方法来进行部分排序；到目前为止我还没有找到一个。

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

原文由 NPE 发布，翻译遵循 CC BY-SA 4.0 许可协议