如何在numpy矩阵的每一行获得N个最大值的索引矩阵？

Question

相关的How to get indices of N maximum values in a numpy array?，我有一个numpy的矩阵a，我想以产生其我 -th行是a的我第行的顶部Ñ元素的列索引的阵列。

继得票最多的答案的链接的问题，适应它为阵，这里是我迄今（使用ñ = 4）：

>>> a 
array([[9, 4, 4, 3, 3, 9, 0, 4, 6, 0], 
     [3, 4, 6, 9, 5, 7, 1, 2, 8, 4]]) 
>>> ind=np.argpartition(a,-4)[:,-4:] 
>>> ind 
array([[1, 5, 8, 0], 
     [2, 3, 8, 5]]) 
>>> rows=np.transpose([np.arange(a.shape[0])]) 
>>> rows 
array([[0], 
     [1]]) 
>>> ind_sorted = ind[rows,np.argsort(a[rows,ind])] 
>>> ind_sorted 
array([[1, 8, 5, 0], 
     [2, 5, 8, 3]]) 

这工作，但似乎是不是很（python）的。我确信有一个更好的方法来做不需要虚拟数组的索引。有什么建议么？

Answer 1

切片按行的顺序索引的最后四个元素似乎是工作：

a.argsort(axis = 1)[:, -4:] 

# array([[7, 8, 0, 5], 
#  [2, 5, 8, 3]]) 

领带方法没有限定，因此会出现的顺序1和7之间的一些差异以及0和5。