MATLAB: extract values from 3d matrix at given row and column indcies using sub2ind 3d - arrays

Question

I have 3d matrix A that has my data. At multiple locations defined by row and column indcies as shown by matrix row_col_idx I want to extract all data along the third dimension as shown below:
A = cat(3,[1:3;4:6], [7:9;10:12],[13:15;16:18],[19:21;22:24]) %matrix(2,3,4)
row_col_idx=[1 1;1 2; 2 3];
idx = sub2ind(size(A(:,:,1)), row_col_idx(:,1),row_col_idx(:,2));
out=nan(size(A,3),size(row_col_idx,1));
for k=1:size(A,3)
temp=A(:,:,k);
out(k,:)=temp(idx);
end
out
The output of this code is as follows:
A(:,:,1) =
1 2 3
4 5 6
A(:,:,2) =
7 8 9
10 11 12
A(:,:,3) =
13 14 15
16 17 18
A(:,:,4) =
19 20 21
22 23 24
out =
1 2 6
7 8 12
13 14 18
19 20 24
The output is as expected. However, the actual A and row_col_idx are huge, so this code is computationally expensive. Is there away to vertorize this code to avoid the loop and the temp matrix?

Answer 1

This can be vectorized using linear indexing and implicit expansion:
out = A( row_col_idx(:,1) + ...
(row_col_idx(:,2)-1)*size(A,1) + ...
(0:size(A,1)*size(A,2):numel(A)-1) ).';
The above builds an indexing matrix as large as the output. If this is unacceptable due to memory limiations, it can be avoided by reshaping A:
sz = size(A); % store size A
A = reshape(A, [], sz(3)); % collapse first two dimensions
out = A(row_col_idx(:,1) + (row_col_idx(:,2)-1)*sz(1),:).'; % linear indexing along
% first two dims of A
A = reshape(A, sz); % reshape back A, if needed

Answer 2

A more efficient method is using the entries of the row_col_idx vector for selecting the elements from A. I have compared the two methods for a large matrix, and as you can see the calculation is much faster.
For the A given in the question, it gives the same output
A = rand([2,3,10000000]);
row_col_idx=[1 1;1 2; 2 3];
idx = sub2ind(size(A(:,:,1)), row_col_idx(:,1),row_col_idx(:,2));
out=nan(size(A,3),size(row_col_idx,1));
tic;
for k=1:size(A,3)
temp=A(:,:,k);
out(k,:)=temp(idx);
end
time1 = toc;
%% More efficient method:
out2 = nan(size(A,3),size(row_col_idx,1));
tic;
for jj = 1:size(row_col_idx,1)
out2(:,jj) = [A(row_col_idx(jj,1),row_col_idx(jj,2),:)];
end
time2 = toc;
fprintf('Time calculation 1: %d\n',time1);
fprintf('Time calculation 2: %d\n',time2);
Gives as output:
Time calculation 1: 1.954714e+01
Time calculation 2: 2.998120e-01