问题描述
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)
将返回索引[4, 3, 1]
对应元素[5, 4, 3]
.
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 个最大值的索引?的文章就介绍到这了,希望我们推荐的答案对大家有所帮助,也希望大家多多支持!