extract 2D matrix from 3D matrix and array of indi

Question

I cannot find out a satisfying answer. If A is a 3D matrix of size (m,n,k), Z is a 2D matrix of size mxn (integers with values between 1 and k), I want to extract S defined like this:

for i=1:m
    for j=n
        S(i,j) = A(i,j,Z(i,j));
    end
end

Is there an efficient (vectorized) way to do this?

Thank you in advance

Answer 1

You can do it using linear indexing as follows:

S = reshape(A((1:m*n).' + m*n*(Z(:)-1)), m, n);