How to create 2D matrix from 3D without a loop?

I have an n x m x 3 matrix A and an n x m matrix B. B contains the values 1, 2, or 3, referring to the three different layers in A. I want to create another n x m matrix C where each point is the value of A from the layer specified by B. For instance, if A(1,1,:) = [8 6 5], B(1,1) = 2, then C(1,1) = 6, because B at that point specifies taking from the second layer. I know how to do this with a loop, but I was wondering if there was an easier, vectorized way to do so. Thanks!

Answer 1

[Andrei Bobrov]

[n,m,k] = size(A);
[i,j] = ndgrid(1:n,1:m);
C = A(sub2ind([n,m,k],i,j,B);

[Humza Khan]

Thanks so much!

Answer 2

[Thomas Satterly]

Here's one way to do it with vectorized code. Calculating the column index list can probably be done better, but it gets the idea across

A = rand(4, 3, 6);

B = [1 2 3; 2 3 4; 3 4 5; 4 5 6];

% List of row indicies
rows = repmat(1:size(B, 1), 1, size(B, 2));

% List of column indicies, probably a better way to do this
cols = repmat(1:size(B, 2), 1, size(B, 1));
cols = reshape(reshape(cols, [size(B, 2), size(B, 1)])', 1, numel(cols));

% Convert row/col to index
inds = sub2ind(size(A), rows, cols, B(:)');

% Get the new matrix with vectorized code
C_vectorized = reshape(A(inds), size(B));

% Get new matrix with loops
C_looped = zeros(size(B));
for i = 1:size(B, 1)
    for j = 1:size(B, 2)
        C_looped(i, j) = A(i, j, B(i, j));
    end
end

% Make sure they're the same
matched = C_vectorized == C_looped;
assert(all(matched(:)), 'Matricies do not match\n');

[Stephen23]

Simpler way to make sure that they are the same:

isequal(C_vectorized,C_looped)