Question

我有一个NxMxK矩阵A

A = [1 2 1; 1 1 2];
A = cat(3, A, [3 3 3; 3 3 3])

A（：，：，1）=
 1     2     1
 1     1     2
A（：，：，2）=
 3     3     3
 3     3     3

我想创建一个YxK 2D矩阵B，其中K是A(:,:,1)==2的元素数：

k=0;
for ii=1:size(A,1)
   for jj=1:size(A,2)
      if A(ii,jj)==2
         k=k+1;
         B(k,:) = A(ii,jj,:);
      end
   end
end

有矢量化此代码的方法吗？

我的尝试是找到A(:,:,1)==2的索引，然后尝试选择整个列，但我不知道该怎么做：

inds = find(A(:,:,1)==2)
B = A(inds,:) %this line does not make sense

编辑预先分配B有帮助：

inds=find(A(:,:,1)==2);
B=NaN(numel(inds),size(A,3));

k=0;
for ii=1:size(A,1)
   for jj=1:size(A,2)
      if A(ii,jj)==2
         k=k+1;
         B(k,:) = squeeze(A(ii,jj,:));
      end
   end
end

但仍然没有向量化。

Answer 1

您可以将reshape矩阵A转换为（N * M）xK 2D矩阵。

A = [1 2 1; 1 1 2];
A = cat(3, A, [3 3 3; 3 3 3]);

A_ = reshape(A,numel(A(:,:,1)),size(A,3));
B = A_(A_(:,1)==2,:);

Answer 2

您第一次进行矢量化的尝试几乎是正确的。只是不使用find，而是使用逻辑矩阵进行索引。

inds = A(:,:,1)==2;

inds矩阵是2D而非3D，因此我们使用repmat沿第3维重复其值：

K = size(A,3);
inds = repmat(inds,1,1,K);  % or simply cat(3,inds,inds) if K==2
B = A(inds);

结果是大小为Y*K的列向量，而不是大小为Y x K的矩阵，我们可以使用reshape来解决这个问题：

B = reshape(B,[],K);

我猜这个答案类似于Anthony's，除了索引和重塑是相反的。直到我写下来之后，我才真正注意到它的相似性。我想安东尼的也要矮一些。：/

根据“表面”值提取3D矩阵的列-矢量化

2 个答案: