I'm struggling to speed up the code inside the version1 function. The only way to vectorize it that I can seem to figure out is version2, but that actually makes the function slower. n1 = 400; n2 = 3; n3 = 7; n4 = 2; A = rand(n1, n2, n3); B = randi(n2, [n4, n1, n3]); disp(timeit(@() version1(A, B, C, n1, n2, n3, n4))); disp(timeit(@() version2(A, B, C, n1, n2, n3, n4))); function C = version1(A, B, n1, n2, n3, n4) C = zeros(n4, n1, n3); for ii = 1:n3 for jj = 1:n1 for kk = 1:n4 C(kk, jj, ii) = A( ... jj, ... B(kk, jj, ii), ... ii ... ); end end end end function C = version2(A, B, n1, n2, n3, n4) C = zeros(n4, n1, n3); for ii = 1:n3 for jj = 1:n1 C(:, jj, ii) = A( ... jj, ... B(:, jj, ii), ... ii ... ); end end end
Kshitij Singh answered .
2025-11-20
[K,J,I]=ndgrid(1:n4,1:n1,1:n3); %Recycle this, if possible Bvals=B(sub2ind(size(B),K,J,I)); C=A( sub2ind(size(A), J,Bvals,I ));