numpy 2d array max/argmax

Question

I have a numpy matrix:

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

I'd like to get the index of the maximum value in each row along with the value itself. I can get the indices for the maximums using A.argmax(axis=1), in that case I would get:

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

How can I use the 'indices' array to get an array of maximum values for each row in the matrix? Is there any way I can do this more efficiently or in one operation? Is there a function that would return the values along with their row and column coordinates?

Answer 1

You can fancy-index using the indices np.arange(len(A)) on first dimension (since you want a value per row), and your indices (squeezed), which correspond to the index on each row, on the second dimension:

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