较新的 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]
的结果可以看出。 如果您也需要它,请稍后对它们进行排序:
#argsort返回从小到大排序后的索引,再进行倒序输出,得到从大到小的索引排序
ind[np.argsort(a[ind])[::-1]]
array([1, 8, 5, 0])