CODEWITHSUNDEEP
matlab
user1581390
user1581390

Reputation: 2018

How to vectorize this operation in Matlab?

The code:

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

How do I write this in vector form, that is using some built it function? I'm basically selecting with the index each point.

Upvotes: 1

Views: 117

Answers (2)

Luis Mendo
Luis Mendo

Reputation: 112769

It can be done as follows:

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

The trick is to build a linear index (ind variable) that corresponds to the indexing (i,j,k,index(i,j,k)), vectorized. I do it with two calls to bsxfun. The key is to remember that linear indices run along columns first, then along rows, then along third-dim slices.

The standard way would be using ndgrid and then sub2ind (see @rayryeng`s answer), but that requires a little more memory.

Upvotes: 1

rayryeng
rayryeng

Reputation: 104555

OK, what this looks like for your output array of B is that for each location in this output, the corresponding location in A has a D-element vector and you want to choose which element to select from this D-element vector based on the value stored in index, which is of the same size as B.

We can achieve this using a combination of meshgrid and sub2ind. First use meshgrid to generate a grid of 3D co-ordinates, which will be the same size as your matrix B. After, use sub2ind to determine a set of linear indices to access the fourth dimension of A. After this, we simply do a straight assignment using the output of sub2ind, which uses a combination of the output of meshgrid and index.

Therefore, try something like this:

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

Upvotes: 1

Related Questions