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 would like to get the index of the maximum value in each row along with the value itself. I can get indexes for maxima using A.argmax (axis = 1), in which case I would get:

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

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

Answer 1

You can index indexes using np.arange(len(A)) indexes in the first dimension (since you need a value for each row), and your indexes (compressed) that correspond to the index for each row in the second dimension:

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