Find maximum of each row in a numpy array and the corresponding element in another array of the same size

Question

I am new to Python and still cannot call myself a Python programmer. Speaking of that, please bear with me if my question does not make any sense.

Question:

I have two numpy arrays of the same size, e.g. A and B where A.shape equals B.shape and they both equal (5,1000), and I want to find the maximum value of each row in A and the corresponding element of that in B. For instance, if in fourth row of A, maximum element index is 104 then I would like to find the 104th element of fourth row in array B and the same for the rest of the rows.

I know I can do it by looping over the rows but I was wondering if there was a more elegant way of doing it. For example, if I were to do it in MATLAB I would write the following code:

B(bsxfun(@eq,A,max(A,[],2)))

Any help that guides me through the right direction would be much appreciated.

Answer 1

Here's the numpy idiom for doing the same thing:

b[np.arange(len(a)), np.argmax(a, axis=1)]

For example:

>>> a = np.array([
    [1, 2, 0],
    [2, 1, 0],
    [0, 1, 2]
    ])
>>> b = np.array([
    [1, 2, 3],
    [1, 2, 3],
    [1, 2, 3]
    ])
>>> b[np.arange(len(a)), np.argmax(a, axis=1)]
array([2, 1, 3])

Answer 2

Being a bsxfun lover, it's great to see people trying to replicate the same functionality to other programming languages. Now, bsxfun is basically a broadcasting mechanism, which exists in NumPy as well. In NumPy, it is achieved by creating singleton dimensions with np.newaxis or simply None.

Back to the question in context, an equivalent broadcasting based solution could be implemented as shown as a sample run -

In [128]: A
Out[128]: 
array([[40, 63, 67, 65, 19],
       [85, 55, 66, 92, 88],
       [50,  1, 23,  6, 59],
       [67, 55, 46, 78,  3]])

In [129]: B
Out[129]: 
array([[78, 63, 45, 34, 81],
       [ 5, 38, 28, 61, 66],
       [ 3, 65, 16, 25, 32],
       [72,  1, 31, 75,  6]])

In [130]: B[A == A.max(axis=1)[:,None]]
Out[130]: array([45, 61, 32, 75])