pick elements from a 3d array, based on an indexing matrix

I have a 3d array A with nr rows, nc columns, and np pages, and an nr x nc indexing matrix I with integer elements between 1 and np. I want to create a matrix B with elements picked from A, where for each position (r,c) the element is picked from the page of A given by the value of I at that position.

A straightforward way to do this would be the double loop:

B = zeros(nr,nc);
for r = 1:nr
    for c = 1:nc
        p = I(r,c);
        B(r,c) = A(r,c,p);
    end
end

Is there a way to accelerate this by vectorization for large arrays?

Answer 1

[Stephen23]

https://www.mathworks.com/help/matlab/ref/sub2ind.html

% fake data:
nr = 5;
nc = 7;
np = 3;
A = randi(9,nr,nc,np);
I = randi(np,nr,nc)

I = 5×7

1 1 2 1 2 2 3 1 3 3 2 3 3 1 1 2 3 2 3 1 1 2 2 2 1 3 2 2 1 1 2 3 3 1 1

% your looped approach:
B = zeros(nr,nc);
for r = 1:nr
    for c = 1:nc
        p = I(r,c);
        B(r,c) = A(r,c,p);
    end
end
B

B = 5×7

7 1 6 2 3 7 6 6 4 3 1 5 5 4 4 1 3 2 2 2 6 4 3 8 2 1 7 9 9 9 6 2 5 1 6

% SUB2IND approach:
[xr,xc] = ndgrid(1:nr,1:nc);
X = sub2ind(size(A),xr,xc,I);
C = A(X)

C = 5×7

7 1 6 2 3 7 6 6 4 3 1 5 5 4 4 1 3 2 2 2 6 4 3 8 2 1 7 9 9 9 6 2 5 1 6

[Bernhard]

Thank you, Stephen, that is a nice gimmick. It took about 60% of the time needed for the looped approach.

Bernhard

[Stephen23]

@Bernhard: if the speed of this operation is very critical to your algorithm, then take a look inside the sub2ind file. The sub2ind function is generic, for any number of dimensions, but with a little effort you can write your own version that is hardcoded for the number of dimensions of your data and that does not require ndgrid beforehand. Whether it is worth the effort depends on your situation.