您好, 欢迎来到 !    登录 | 注册 | | 设为首页 | 收藏本站

Python-如何获取NumPy数组中N个最大值的索引?

Python-如何获取NumPy数组中N个最大值的索引?

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

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提供了一种进行部分排序的内置方法。到目前为止,我还没有找到一个

如果此解决方案太慢(尤其是对于小型解决方案n),则可能值得考虑使用Cython进行编码。

较新的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 log k)时间。

python 2022/1/1 18:28:52 有189人围观

撰写回答


你尚未登录,登录后可以

和开发者交流问题的细节

关注并接收问题和回答的更新提醒

参与内容的编辑和改进,让解决方法与时俱进

请先登录

推荐问题


联系我
置顶