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

Question

如何在NumPy数组中获得N个最大值的索引？NumPy提出了一种方法，通过np.argmax.我想要一个类似的东西，但是返回N最大值。例如，如果我有一个数组，[1, 3, 2, 4, 5], function(array, n=3)会返回指数[4, 3, 1]与元素对应的[5, 4, 3].

Answer 1

我想出的最简单的方法是：

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

这涉及到一个完整的数组类型。我想知道numpy提供一个内置的方式来进行部分排序；到目前为止，我还没有找到一个。

如果这个解决方案被证明太慢(特别是对于小型的)n)，可能值得考虑一下在Cython.

Answer 2

更新的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])>>> a[ind]array([4, 9, 6, 9])

不像argsort，该函数在最坏的情况下以线性时间运行，但返回的索引没有排序，从计算结果可以看出。a[ind]..如果你也需要的话，把它们整理一下：

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

为了得到最高层-k按排序顺序排列的元素以这种方式取O(n + k原木k)时间。

Answer 3

用途：

>>> import heapq>>> import numpy>>> a = numpy.array([1, 3, 2, 4, 5])>>> heapq.nlargest(3, range(len(a)), a.take)[4, 3, 1]

对于常规Python列表：

>>> a = [1, 3, 2, 4, 5]>>> heapq.nlargest(3, range(len(a)), a.__getitem__)[4, 3, 1]

如果使用Python 2，请使用xrange而不是range.

资料来源：堆队列算法