如何获得 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) 将返回索引[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 个最大值的索引?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！