I am using Matlab R2014a.
I have a 3-dimensional M x N x M matrix A. I would like a vectorized way to extract a 2 dimensional matrix B from it, such that for each i,j I
You can employ the best tool for vectorization, bsxfun here -
B = A(bsxfun(@plus,[1:M]',M*(0:N-1)) + M*N*(g-1))
Step #1: Calculate the indices corresponding to the first two dimensions (rows and columns) of A
-
bsxfun(@plus,[1:M]',M*(0:N-1))
Step #2: Add the offset needed to include the dim-3 indices being supplied by g
and index into A with those indices to get our desired output -
A(bsxfun(@plus,[1:M]',M*(0:N-1)) + M*N*(g-1))
Here's a quick benchmark test to compare this bsxfun
based approach against the ndgrid + sub2ind
based solution as presented in Luis's solution with M
and N
as 100
.
The benchmarking code using tic-toc
would look something like this -
M = 100;
N = 100;
A = rand(M,N,M);
g = randi(M,M,N);
num_runs = 5000; %// Number of iterations to run each approach
%// Warm up tic/toc.
for k = 1:50000
tic(); elapsed = toc();
end
disp('-------------------- With BSXFUN')
tic
for iter = 1:num_runs
B1 = A(bsxfun(@plus,[1:M]',M*(0:N-1)) + M*N*(g-1)); %//'
end
toc, clear B1
disp('-------------------- With NDGRID + SUB2IND')
tic
for iter = 1:num_runs
[ii, jj] = ndgrid(1:M, 1:N);
B2 = A(sub2ind([M N M], ii, jj, g));
end
toc
Here's the runtime results -
-------------------- With BSXFUN
Elapsed time is 2.090230 seconds.
-------------------- With NDGRID + SUB2IND
Elapsed time is 4.133219 seconds.
As you can see bsxfun
based approach works really well, both as a vectorized approach and good with performance too.
Why is bsxfun
better here -
bsxfun
does replication of offsetted elements and adding them, both on-the-fly.
In the other solution, ndgrid
internally makes two function calls to repmat
, thus incurring the function call overheads. At the next step, sub2ind
spends time in adding the offsets to get the linear indices, bringing in another function call overhead.
Try using sub2ind. This assumes g
is defined as an M
xN
matrix with possible values 1
, ..., M
:
[ii, jj] = ndgrid(1:M, 1:N);
B = A(sub2ind([M N M], ii, jj, g));