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]