Matlab code performance improvement

Question

I have a matrix A (MxN) and need to create a matrix B (MxNxN) by using A such that B(:,1,1) = A(:,1), B(:,2,2) = A(:,2),..., B(:,N,N) = A(:,N). Currently I use,

B = zeros(size(A,1), size(A,2), size(A,2));
for i=1:size(B,3)
    B(:,i,i) = A(:,i);
end

Is it possible to use repmat or any other method to make this code run faster than it does now?

Answer 1

It can be done using "partial" linear indexing on the last two dimensions of B, as follows:

[M, N] = size(A);
B = zeros(M, N, N);
B(:, 1:N+1:N^2) = A;