arrays - Matlab代码性能提升

Question

我有一个矩阵 A (MxN) 并且需要使用 A 创建一个矩阵 B (MxNxN) 使得 B(:,1,1) = A(:,1), B(:,2,2) = A(:,2),..., B(:,N,N) = A(:,N)。目前我在用，

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

是否可以使用 repmat 或任何其他方法使这段代码运行得比现在更快？

Answer 1

可以使用 "partial" linear indexing 来完成在B的最后两个维度上，如下:

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