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

发布于 2021-02-02 23:18:57

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

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

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

关注者
0
被浏览
141
1 个回答
  • 面试哥
    面试哥 2021-02-02
    为面试而生,有面试问题,就找面试哥。

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

    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进行编码。



  • 面试哥
    面试哥 2021-02-02
    为面试而生,有面试问题,就找面试哥。

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



知识点
面圈网VIP题库

面圈网VIP题库全新上线,海量真题题库资源。 90大类考试,超10万份考试真题开放下载啦

去下载看看