Using index of 3d array from max option to access elements from another 3d array

Here is an example of code that executes what I'd like to do, but is very slow for large arrays and can likely be written much more simply and faster. I have an array A that I take the max over the rows, where A is a 3d array. I'd then like to use the indices from the max operator to index into another array that is the same size as B and collect those elements in an array where the size of the final array is [1,size(A,2),size(A,3)]. Any help would be greatly appreciated. Thank you!

     A = rand(3,2,10);
     [~,Aind] = max(A,[],1);

     B = rand(3,2,10);

    tmp3 = zeros(1,size(A,2),size(A,3));
    for j = 1:size(A,2)
        tmp1 = squeeze(Aind(:,j,:))+ 3*((0:1:10-1)');
        tmp2 = squeeze(B(:,j,:));
        tmp3(1,j,:) = tmp2(tmp1);
    end

Answer 1

[Guillaume]

I always find it difficult to think in 3D, but the following should do what you want. In any case, sub2ind (and ind2sub) is your friend.

[~, Aind] = max(A, [], 1);
tmp3 = B(sub2ind(size(B), Aind, ...
                          repmat(1:size(A, 2), 1, 1, size(A, 3)), ...
                          repmat(permute(1:size(A, 3), [1 3 2]), 1, size(A, 2), 1)))

A little squeeze beforehand would make it slightly more readable:

tmp3 = B(sub2ind(size(B), squeeze(Aind), ...
                          repmat((1:size(A, 2))', 1, size(A, 3)), ...
                          repmat(1:size(A, 3), size(A, 2), 1)))

[jonathan williams]

Thank you Guillaume!

[hamed amini]

a more generic solution and could be extended to more dimensions:

[a2, a3] = ndgrid([1:size(A,2)],[1:size(A,3)]);

tmp3 = B(sub2ind(size(B), squeeze(Aind), a2, a3)