Vectorized nD diag loop

Hello everyone,

I have the following loop which I would want to have vectorized.

I know the diag() command cannot be used in nD arrays (n>2), so is there an alternative way to vectorize the loop?

n = 10; m = 15;
A = rand(4,n,m);
for i = 1:n
    for j = 1:m
B(:,:,i,j) = diag(A(:,i,j));
    end
end

Thanks in advance!

Answer 1

[Chunru]

n = 400; m = 500;
A = rand(4,n,m);
B = zeros(4,4,n,m);   % initialize to speed up
tic
for i = 1:n
    for j = 1:m
        B(:,:,i,j) = diag(A(:,i,j));
    end
end
toc

Elapsed time is 0.351235 seconds.

% Vectorized
tic
C = zeros(4*4,n*m);
C(1:5:16, :) = A(1:4, :);
C =reshape(C, [4, 4, n, m]);
toc

Elapsed time is 0.015569 seconds.

isequal(B, C)

ans = logical
   1

[Javier Fernandez]

This is what I was looking for! Thank you very much for your kind help.