2D numpy argsort索引在原始矩陣中使用時返回3D

Question

我正嘗試使用argsort從矩陣中的每一行獲取前2個值。索引工作正常，因爲argsort正在返回正確的值。但是，當我將argsort結果作爲索引時，它會返回3維結果。

例如：

test_mat = np.matrix([[0 for i in range(5)] for j in range(5)]) 
for i in range(5): 
    for j in range(5): 
     test_mat[i, j] = i * j 
test_mat[range(2,3)] = test_mat[range(2,3)] * -1 

last_two = range(-1, -3, -1) 
index = np.argsort(test_mat, axis=1) 
index = index[:, last_k] 

這給：

index.shape 
Out[402]: (5L, 5L) 

test_mat[index].shape 
Out[403]: (5L, 5L, 5L) 

Python是新的給我，我發現索引，甚至閱讀的各種陣列手冊後，在一般非常混亂。我花更多時間試圖從對象中獲得正確的值，而不是實際解決問題。我歡迎有關正確瞭解正在發生的事情的任何提示。謝謝。

Answer 1

您可以使用linear indexing來解決你的情況下，像這樣 -

# Say A is your 2D input array 

# Get sort indices for the top 2 values in each row 
idx = A.argsort(1)[:,::-1][:,:2] 

# Get row offset numbers 
row_offset = A.shape[1]*np.arange(A.shape[0])[:,None] 

# Add row offsets with top2 sort indices giving us linear indices of 
# top 2 elements in each row. Index into input array with those for output. 
out = np.take(A, idx + row_offset) 

這裏有一個一步一步的樣品運行 -

In [88]: A 
Out[88]: 
array([[34, 45, 16, 20, 24], 
     [37, 13, 49, 37, 21], 
     [42, 36, 35, 24, 18], 
     [26, 28, 21, 13, 44]]) 

In [89]: idx = A.argsort(1)[:,::-1][:,:2] 

In [90]: idx 
Out[90]: 
array([[1, 0], 
     [2, 3], 
     [0, 1], 
     [4, 1]]) 

In [91]: row_offset = A.shape[1]*np.arange(A.shape[0])[:,None] 

In [92]: row_offset 
Out[92]: 
array([[ 0], 
     [ 5], 
     [10], 
     [15]]) 

In [93]: np.take(A, idx + row_offset) 
Out[93]: 
array([[45, 34], 
     [49, 37], 
     [42, 36], 
     [44, 28]]) 

---

你可以直接得到前2從第二個軸和一些slicing排序每個行的值，像這樣 -

out = np.sort(A,1)[:,:-3:-1]