如何获取NumPy数组中N个最大值的索引? [英] How do I get indices of N maximum values in a NumPy array?

问题描述

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

NumPy proposes a way to get the index of the maximum value of an array via np.argmax.

我想要类似的事情，但是返回N最大值的索引.

I would like a similar thing, but returning the indexes of the N maximum values.

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

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

推荐答案

我能想到的最简单的方法是:

The simplest I've been able to come up with is:

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是否提供了一种进行部分排序的内置方法.到目前为止，我还找不到一个.

This involves a complete sort of the array. I wonder if numpy provides a built-in way to do a partial sort; so far I haven't been able to find one.

如果该解决方案太慢(尤其是对于小型n而言)，则可能值得在 Cython .

If this solution turns out to be too slow (especially for small n), it may be worth looking at coding something up in Cython.

这篇关于如何获取NumPy数组中N个最大值的索引?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！