Vectorizing ther higher dimensions in nested for l

2019-08-07 06:22发布

问题:

I have a 5D matrix A, and I need to multiply the 3rd-5th dimensions with a vector. For example, see the following sample code:

A=rand(50,50,10,8,6);
B=rand(10,1);
C=rand(8,1);
D=rand(6,1);

for i=1:size(A,3)
    for j=1:size(A,4)
        for K=1:size(A,5)
            A(:,:,i,j,K)=A(:,:,i,j,K)*B(i)*C(j)*D(K);
        end
    end
end

I wonder if there's a better \ vectorized \ faster way to do this?

回答1:

Firstly, as a note, these days in Matlab, with JIT compilation, vectorised code is not necessarily faster/better. For big problems the memory usage in particular can cause performance issues.

Nevertheless, here is a vectorised solution that seems to give the same results as your code:

A=rand(3,4,5,6,7);
B=rand(5,1);
C=rand(6,1);
D=rand(7,1);   

s=size(A);
[b,c,d]=ndgrid(B,C,D);
F=b.*c.*d;
G=zeros(1,1,prod(s(3:5)));
G(1,1,:)=F(:);
A=reshape(A,s(1),s(2),[]);
A=bsxfun(@times,A,G);
A=reshape(A,s);

EDIT: An alternate solution:

A=bsxfun(@times,A,permute(B,[2 3 1]));
A=bsxfun(@times,A,permute(C,[2 3 4 1]));
A=bsxfun(@times,A,permute(D,[2 3 4 5 1]));