Arg max without using a loop

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

Answer 1

[Adam]

[ 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.

Answer 2

[Luca Gagliardone]

Thanks Adam, I see your point.

It still seems to me impossible that Matlab does not provide a more direct way without recurring to linear indexing (which is also quite RAM-expensive with ndgrid). I will wait for a few days before accepting definitely your answer, hoping for an easier solution!

Thanks again for your help.

Luca

[Adam]

Well, everything is linear indexing under the hood anyway.

[Luca Gagliardone]

Well, that's true!