How to just calculate the diagonal of a matrix pro

Question

I have two matrix A and B, so what's the fastest way to just calculate diag(A%*%B), i.e., the inner-product of the ith row of A and ith column of B, and the inner-product of other terms are not concerned.

supplement: A and B have large row and column numbers respectively.

Answer 1

This can be done without full matrix multiplication, using just multiplication of matrix elements.

We need to multiply rows of A by the matching columns of B and sum the elements. Rows of A are columns of t(A), which we multiply element-wise by B and sum the columns.

In other words: colSums(t(A) * B)

Testing the code we first create sample data:

n = 5
m = 10000;

A = matrix(runif(n*m), n, m);
B = matrix(runif(n*m), m, n);

Your code:

diag(A %*% B)
# [1] 2492.198 2474.869 2459.881 2509.018 2477.591

Direct calculation without matrix multiplication:

colSums(t(A) * B)
# [1] 2492.198 2474.869 2459.881 2509.018 2477.591

The results are the same.