matlab中单个循环的矢量化（乘法然后加法）

Question

我有一个 nX2 矩阵 A 和一个 3D 矩阵 K。我想进行元素乘法，指定由 A 中的每个行向量指定的 K 的第三维中的 2 个索引，并对它们求和。

例如当 n=2 时的一个简化示例，

A=[1 2;3 4];%2X2 matrix
K=unifrnd(0.1,0.1,2,2,4);%just random 3D matrix
L=zeros(2,2);%save result to here
for t=1:2
    L=L+prod(K(:,:,A(t,:)),3);
end

在这种情况下我可以摆脱 for 循环吗？

Answer 1

怎么样？

B = A.'; %'
L = squeeze(sum(prod(...
                  reshape(permute(K(:,:,B(:)),[3 1 2]),2,[],size(K,1),size(K,2)),...
                  1),...
            2));

虽然你的测试用例太简单了，所以我不能完全确定它是否正确。

我们的想法是，我们首先以列优先顺序获取A 中的所有索引，然后重塑K 的元素，使得前两个维度的大小为[2, n]，后两个维度的大小是 K 的原始 2。然后我们取乘积，然后是沿必要维度求和，最后得到一个矩阵，该矩阵必须被压缩才能得到一个二维矩阵。

---

使用更多信息的测试用例：

K = rand(2,3,4);
A = randi(4,4,2);
L = zeros(2,3);%save result to here
for t=1:size(A,1)
    L = L+prod(K(:,:,A(t,:)),3);
end

B = A.'; %'
L2 = squeeze(sum(prod(reshape(permute(K(:,:,B(:)),[3 1 2]),2,[],size(K,1),size(K,2)),1),2));

然后

>> isequal(L,L2)
ans =
     1

Answer 2

有一些reshaping magic -

%// Get sizes
[m1,n1,r1] = size(K);
[m2,n2] = size(A);

%// Index into 3rd dim of K; perform reductions and reshape back
Lout = reshape(sum(prod(reshape(K(:,:,A'),[],n2,m2),2),3),m1,n1);

解释：

1. 使用A 的转置版本对K 的第三维进行索引（转置是因为我们使用A 的行进行索引）。
2. 执行prod() 和sum() 操作。
3. 最后重新整形为与 K 相同的形状，但没有在前面的缩减步骤中删除的第三维。