如何在 Matlab 中对这个操作进行矢量化？

Question

代码：

A = NxMxLxD matrix
B = NxMxL matrix
index = NxMxL matrix containing values between 1 and D

for i=1:N
 for j=1:M
  for k=1:L
   B(i,j,k) = A(i,j,k,index(i,j,k));
  end
 end
end

如何以矢量形式编写它，即使用一些内置的函数？我基本上是用每个点的索引来选择的。

Answer 1

好的，对于您的B 输出数组，这看起来是对于此输出中的每个位置，A 中的相应位置有一个D-元素向量，并且您要选择要选择的元素来自这个D-element 向量，基于存储在index 中的值，其大小与B 相同。

我们可以使用meshgrid 和sub2ind 的组合来实现这一点。首先使用meshgrid 生成一个3D 坐标网格，其大小将与您的矩阵B 相同。之后，使用sub2ind 确定一组线性索引以访问A 的第四维。在此之后，我们只需使用sub2ind 的输出进行直接赋值，它使用meshgrid 和index 的输出的组合。

因此，请尝试以下操作：

[cols,rows,dim] = meshgrid(1:size(A,2), 1:size(A,1), 1:size(A,3));
ind = sub2ind(size(A), rows, cols, dim, index);
B = A(ind);

Answer 2

可以这样做：

ind = bsxfun(@plus, bsxfun(@plus, ...
      (1:N).', N*(0:M-1)), N*M*permute(0:L-1, [3 1 2])) + N*M*L*(index-1);
B = A(ind);

诀窍是构建一个linear index（ind 变量）对应于索引(i,j,k,index(i,j,k))，矢量化。我通过两次调用bsxfun 来做到这一点。关键是要记住线性索引首先沿列运行，然后沿行运行，然后沿第三维度切片运行。

标准方法是使用ndgrid，然后使用sub2ind（请参阅@rayryeng`s answer），但这需要更多内存。