Value of B at the maximum of A

My problem is quite simple. I have two arrays A and B of identical size and having N dimensions.

I want to compute the max of A along its Nth dimension and to get the corresponding values in array B.

I see several complicated ways to to it, but I guess that there must be an elegant and efficient one I am missing.

Thanks a lot for your help,

-------

Edit (cf. comments) :

A non-vectorized version of I want is:

A = randi(99, 4, 5, 3) ;
B = randi(99, 4, 5, 3) ;
[Amax,Idx] = max(A,[],3) ;

for i=1:4 ;
   for j=1:5 ;  
         C(i,j) = B(i,j,Idx(i,j)) ;
   end
end

[Star Strider]

Three hours ago that would have been helpful.

Answer 1

[Sylvain Catherine]

Here is a way to do it. I am not sure it's the best way however...

size1 = 4 ;
size2 = 5 ;
size3 = 3 ;

A = randi(99, size1, size2, size3) ;
B = randi(99, size1, size2, size3) ;
[Amax,Idx] = max(A,[],3) ;

IDX = 1:size1*size2 ;
IDX = reshape(IDX,size1,5) + (Idx-1)*size1*size2;

C = B(IDX) ;

Answer 2

[Star Strider]

Unless I don’t understand what you want, you can use the max function with two outputs.

For example, to get the maximum along the third dimension of a 3-dimensional matrix, and then get the corresponding elements from another matrix of the same size this works:

A = randi(99, 4, 5, 3)
B = randi(99, 4, 5, 3)

[Amax,Idx] = max(A,[],3)
C = B(:,:,Idx)                          % Desired Result

[Walter Roberson]

Ah! I didn't think of just using the matrix itself as an index.

[Sylvain Catherine]

Thanks Star Strider, but your code produces a matrix C of size 4x5x20, while I would like C and Amax to have identical sizes (4x5).

A non-vectorized version of I want is:

A = randi(99, 4, 5, 3) ;
B = randi(99, 4, 5, 3) ;
[Amax,Idx] = max(A,[],3) ;

for i=1:4 ;
   for j=1:5 ;  

         C(i,j) = B(i,j,Idx(i,j)) ;

   end
end

[Walter Roberson]

Generalizing...

A = randi(99, 4, 5, 3);
B = randi(99, 4, 5, 3);

N = randi(ndims(A));    %random dimension

[Amax, idx] = max(A, [], N);
idx_list = repmat({':'}, 1, ndims(A));
idx_list{N} = idx;
C = B(idx_list{:});

... Except it does not work. When you use a matrix as index then the ':' for the other dimensions is still respected, so you end up with something that is (for example) 4 x 5 x (4*5) rather than 4 x 5 x 3.

[Star Strider]

The Amax and Idx matrices are themselves both 4x5 (in my example code), so they meet your expectations.

In creating ‘C’, it then ‘pages’ through the dimensions of ‘B’ and creates the 20 dimension ‘C’ matrix by creating a single (4x5) matrix from each index in ‘Idx’. Walter’s analysis is correct, as is the result MATLAB produces by default.

If the result of your ‘non-vectorised’ loop is the result you want, then go with it. You did not specify the result you wanted in your original Question.

[Walter Roberson]

I know a vectorized version is possible; at the moment it is not clear to me that it can be done efficiently.