ndarray每行中N個最大值

Question

我有一個ndarray，其中每行是一個單獨的直方圖。對於每一行，我希望找到前N個值。

我知道一個解決方案的全球前N值（A fast way to find the largest N elements in an numpy array），但我不知道如何獲得每行的前N。

我可以遍歷每一行並應用1D解決方案，但是不應該用numpy廣播來做到這一點嗎？

Answer 1

您可以使用np.partition以與您鏈接的問題相同的方式：排序已經沿着最後一個軸：

In [2]: array([[ 5, 4, 3, 2, 1], 
       [10, 9, 8, 7, 6]]) 
In [3]: b = np.partition(a, 3) # top 3 values from each row 
In [4]: b[:,-3:] 
Out[4]: 
array([[ 3, 4, 5], 
     [ 8, 9, 10]]) 

Answer 2

您可以使用np.argsort沿着行與axis = 1像這樣 -

import numpy as np 

# Find sorted indices for each row 
sorted_row_idx = np.argsort(A, axis=1)[:,A.shape[1]-N::] 

# Setup column indexing array 
col_idx = np.arange(A.shape[0])[:,None] 

# Use the column-row indices to get specific elements from input array. 
# Please note that since the column indexing array isn't of the same shape 
# as the sorted row indices, it will be broadcasted 
out = A[col_idx,sorted_row_idx] 

採樣運行 -

In [417]: A 
Out[417]: 
array([[0, 3, 3, 2, 5], 
     [4, 2, 6, 3, 1], 
     [2, 1, 1, 8, 8], 
     [6, 6, 3, 2, 6]]) 

In [418]: N 
Out[418]: 3 

In [419]: sorted_row_idx = np.argsort(A, axis=1)[:,A.shape[1]-N::] 

In [420]: sorted_row_idx 
Out[420]: 
array([[1, 2, 4], 
     [3, 0, 2], 
     [0, 3, 4], 
     [0, 1, 4]], dtype=int64) 

In [421]: col_idx = np.arange(A.shape[0])[:,None] 

In [422]: col_idx 
Out[422]: 
array([[0], 
     [1], 
     [2], 
     [3]]) 

In [423]: out = A[col_idx,sorted_row_idx] 

In [424]: out 
Out[424]: 
array([[3, 3, 5], 
     [3, 4, 6], 
     [2, 8, 8], 
     [6, 6, 6]]) 

如果你想有降序排列的元素，你可以使用這個額外步 -

In [425]: out[:,::-1] 
Out[425]: 
array([[5, 3, 3], 
     [6, 4, 3], 
     [8, 8, 2], 
     [6, 6, 6]])