Use index of maximum value on a defined axis

Question

I want to extract the values of an array "B" at the same index of the maximum of each line of the matrix "A". to find the index I use the numpy function "numpy.argmax" like this:

>>> A=numpy.asarray([[0,1,6],[3,2,4]]);A
array([[0, 1, 6],
       [3, 2, 4]])
>>> argA=numpy.argmax(A,axis=1);argA
array([2, 2])

The problem is that I don't know how to use "argA" to extract the values in the array "B"

Answer 1

Each entry in argA corresponds to the index position of a maximum value within a corresponding row. The rows are not explicit (due to using axis=1), but correspond to the index for each entry. So you need to add them in to get at the elements you are after.

>>> A[[0,1], argA]
array([6, 4])

So:

>>> B
array([[ 9,  8,  2],
       [ 3,  4,  5]])
>>> B[[0,1], argA] = 84,89
>>> B
array([[ 9,  8, 84],
       [ 3,  4, 89]])

to generalise use:

>>> B[np.arange(A.shape[0]),argA]