numpy 2d array max / argmax

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))進行花式索引（因為您希望每行得到一個值），並且您的索引（壓縮）對應於第二維上每行的索引：

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