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]