numpy 二维数组 最大值/索引

Question

我有一个numpy矩阵：

>>> A = np.matrix('1 2 3; 5 1 6; 9 4 2')
>>> A
matrix([[1, 2, 3],
        [5, 1, 6],
        [9, 4, 2]])

我想要找出每一行最大值的索引和这个最大值本身。我可以用A.argmax(axis=1)来获取每行最大值的索引，这样我会得到：

>>> indices = A.argmax(axis=1)
>>> indices
matrix([[2],
        [2],
        [0]])

我该如何利用这个'indices'数组来获取每一行的最大值数组呢？有没有更高效的方法或者一次性完成这个操作的办法？有没有什么函数可以同时返回最大值及其所在的行和列坐标？

Answer 1

你可以使用索引 np.arange(len(A)) 来对第一维进行 fancy-indexing（花式索引），因为你想要每一行一个值。然后在第二维上使用你的索引（经过压缩），这些索引对应于每一行的具体位置：

A[np.arange(len(A)) , indices.squeeze()]
=> matrix([[3, 6, 9]])