Help with a vectorizing: rearranging a matrix

I was hoping for help with vectorizing the following:

say I have a matrix A and a function F: (i,j) -> j'. I would like to create a new matrix B with the rule: B(i,j) = A(i,F(i, j)).

This can be done with a forloop, but F is very quick to compute and the overhead in running this forloop appears to be a massive bottleneck for me. Is there a better way to write it than the following? I tried arrayfun but this seems to have an even worse overhead?

(Forloop)

B = zeros(size(A));

for i = 1:size(A,1)

for j = 1:size(A,2)

B(i,j) = A(i, F(i,j));

end

end

[James Tursa]

Is the F function vectorized? I.e., if i and j are same-sized vectors, will F return a same-sized vector?

[Dylan Altschuler]

Hi James,

yes, I think so!

Answer 1

[Aditya Patil]

Assuming F itself can be vectorized, first calculate all required values of F.

len = 5;
F = randi(len, [len, len]); % Replace this with vectorized calculation of F
A = rand([len, len]);

Next, generate the indices to be used for A, using sub2ind.

i = 1:len;
j = 1:len;
is = repmat(i, 1, len);
Fs = reshape(F', 1, []);
Aind = sub2ind(size(A), is, Fs);

Then use these indices to update B.

B = A(Aind);

[Dylan Altschuler]

Thank you! I appreciate the help.

I will time this to check if it performs better than array fun