Matlab中嵌套for循环中的高维矢量化

Question

我有一个5D矩阵A，我需要用一个向量乘以第3-5维。例如，请参阅以下示例代码：

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

我想知道是否有更好的\矢量化\更快的方法来做到这一点？

Answer 1

首先，值得注意的是，在Matlab中，使用JIT编译，矢量化代码不一定更快/更好。对于大问题，尤其是内存的使用可能会导致性能问题。

然而，这里有一个矢量化的解决方案，其结果似乎与您的代码相同：

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);

编辑:另一种解决方案：

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]));