使用argsort结果进行numpy索引[重复]

Posted 2021-05-06

这个问题在这里已有答案：

• Get N maximum values and indices along an axis in a NumPy array 3回答

我得到了一个numpy 2D数组，以及与使用argsort获得的前3个元素对应的索引列表。现在，我正在尝试提取与此索引相对应的值，但它无法正常工作。解决方法是什么？

A = array([[0.19334242, 0.9787497 , 0.41453434, 0.35298119, 0.17943745,
        0.63468207, 0.43840688],
       [0.39811914, 0.68040634, 0.7589702 , 0.3573046 , 0.16365397,
        0.86329535, 0.48559053],
       [0.5848541 , 0.54203383, 0.27262654, 0.21979374, 0.06917679,
        0.10586995, 0.57083441],
       [0.76765549, 0.05703751, 0.83383973, 0.71867625, 0.16338699,
        0.85721418, 0.5953548 ]])

np.flip(A.argsort(),axis=1)[:,0:3]
array([[1, 5, 6],
       [5, 2, 1],
       [0, 6, 1],
       [5, 2, 0]])

得到错误

>>> A[np.flip(A.argsort(),axis=1)[:,0:3]]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
IndexError: index 5 is out of bounds for axis 0 with size 4

答案

In [22]: A.ravel()[A.argsort(axis=None)[::-1][:3]]
Out[22]: array([ 0.9787497 ,  0.86329535,  0.85721418])

Explanation

默认情况下，argsort()沿最后一个轴排序。在您的情况下，您希望对数组的扁平版本进行排序，因为您对数组是2D的事实没有任何意义。这是通过将axis=None传递给argsort()而发生的。

由于你获得了1D索引，你还需要访问数组的扁平版本的值，这就是ravel()所做的。

[::-1]反转argsort数组以获得最高值，[:3]获得前3个值。

注意：还有其他可能更有效的方法，但这是我想到的第一件事。