如何在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。