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

Question

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

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

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

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) 时间。