MATLAB: Arg max without using a loop

argmaxmax

I am looking for a method to reproduce the following without using a loop.
The issues come from the fact that the matrices I am using are too big and I could save some space avoiding to create a third matrix C to do the job.
My aim is to find some function f that does something similar to B = f(B(policy)). Is it possible to do it?
A = rand(10,20,30);
B = rand(10,20,30);
[A_hat, policy] = max(A,[],3);
for i = 1:10
for ii = 1:20
C = B(i,ii,policy(i,ii));
end
end
Thanks for the attention.
Sincerely
Luca

Best Answer

[ rowGrid, columnGrid ] = ndgrid( 1:size(A,1), 1:size(A,2) );
idx = sub2ind( size( B ), rowGrid, columnGrid, policy );
C = B( idx );
As far as I am aware you have to convert to linear indices to extract the data as above. I have done exactly this for something I am working on recently. There may be a more efficient way, but not that I am aware of.
I tested one or two values to check they were correct and also looking at
A( idx )
gives a good confidence level too since these numbers are obviously all towards the higher end of the 0-1 range, e.g.
>> stuff = A( idx );
>> mean( stuff(:) )
ans =
           0.969286594713116
Note: This may be memory intensive though as you create two grids of size(A,1) * size(A,2) containing the row and column indices in order to be able to turn all your 3rd dimension indices into linear indices.
Related Question