Composition of two matrices with the same number of columns, one of them is integer-valued

I have two matrices A, B with the same number of columns (d). The entries of B are positive integers. I want to create a matrix C defined by C(i,j)=A(B(i,j),j) and find the products of the rows of C.

I tried the for loop below but it is quite slow. Is there any MATLAB operation to compute this kind of composition faster?

B = zeros(n,d); 
m = max(B(:)); 
A = zeros(m,d);
C = zeros(n,d); 
for ii = 1:n 
    for jj = 1:d 
        C(ii,jj) = A(B(ii,jj),jj); 
    end
end
prod_C = prod(C,2); 

[Jan]

What are typical sizes of n, m and d? Dozens or billions?

[Quang Huy Pham]

These matrices are not big: d is dozens, n is thousands, m is less than 10. But I need to run the code repeatedly millions of times for simulation. The index matrix B is fixed throughout, A changes over time.

Answer 1

[Jan]

n = 1e4;
m = 1e4;
d = 1e3;

B = randi([1, m], n, d);  
A = randn(m, d) * 2;

% Version 1: The original code
tic
C = zeros(n, d); 
for ii = 1:n 
    for jj = 1:d 
        C(ii,jj) = A(B(ii,jj),jj); 
    end
end
prod_C = prod(C,2);
toc

Elapsed time is 0.601715 seconds.

% Version 2: Change the order of loops
tic
C = zeros(n, d); 
for jj = 1:d 
   for ii = 1:n 
       C(ii,jj) = A(B(ii,jj),jj); 
   end
end
prod_C2 = prod(C,2);
toc

Elapsed time is 0.170638 seconds.

% Version 3: Omit the inner loop
tic
C = zeros(n, d); 
for jj = 1:d 
   C(:, jj) = A(B(:, jj), jj); 
end
prod_C3 = prod(C, 2);
toc

Elapsed time is 0.114302 seconds.

% Version 4: Multiply on the fly (no large intermediate array!)
tic
prod_C4 = ones(n, 1); 
for jj = 1:d 
   prod_C4 = prod_C4 .* A(B(:, jj), jj); 
end
toc

Elapsed time is 0.073988 seconds.

10 times faster finally.

A nit-picking improvement of version: 4

% Version 4b: Should not change the speed measureably
tic
prod_C4 = A(B(:, 1), 1);  % Save one multiplication
for jj = 2:d              % Start at jj=2
   prod_C4 = prod_C4 .* A(B(:, jj), jj); 
end
toc

Elapsed time is 0.072217 seconds.

% [EDITED] 4c: Mulitply on the fly in a loop:
tic
prod_C4 = ones(n, 1);
for jj = 1:d
   for ii = 1:n
      prod_C4(ii) = prod_C4(ii) * A(B(ii,jj),jj);
   end
end
toc

Elapsed time is 0.121049 seconds.

% Version 5 - linear indexing, no loop, but large intermediate arrays:
tic;
prod_C5 = prod(A(B + (0:d-1)*m), 2);
toc

Elapsed time is 0.182863 seconds.

% Are all results equal?
isequal(prod_C, prod_C2, prod_C3, prod_C4, prod_C5)

ans = logical
   1

Choose version 4.

[EDITED] In my local Matlab 2018b, version 4c is 30% faster than version 4.

[Quang Huy Pham]

Thank you very much! I tried version 4. It makes my code 5 times faster.

[Jan]

The timings in my local R2018b version differ: Here version 2 is the fastest. But see [EDITED] version 4c for a faster one.