Diagonal part of product of matrices

Question

I am trying to find the best way to compute only the diagonal part of a product of two (or three) matrices. I tried two different approaches: computing the standard matrix product and then extracting the diagonal part; extracting rows and columns from the external matrices and making a Table. The first approach is of course very costly as it computes many more elemenst than needed. And this becomes evident with two matrices

matrA = RandomReal[NormalDistribution[], {3000, 3000}];
matrB = RandomReal[NormalDistribution[], {3000, 3000}];
Timing[vec0 = Diagonal[matrA.matrB];]
Timing[Block[{matrBT = Transpose[matrB]}, vec1 = Table[matrA[[i]].matrBT[[i]], {i, 1, Length[matrA]}]];]

As expected, the first expression takes much longer than the second one. However, I was surprised to find out that with three matrices, this does not seem to be true anymore.

matrA = RandomReal[NormalDistribution[], {3000, 3000}];
matrB = RandomReal[NormalDistribution[], {3000, 3000}];
matrC = RandomReal[NormalDistribution[], {3000, 3000}];
Timing[vec0 = Diagonal[matrA.matrC.matrB];]
Timing[Block[{matrBT = Transpose[matrB]}, vec1 = Table[matrA[[i]].matrC.matrBT[[i]], {i, 1, Length[matrA]}]];]

Now surprisingly the first approach is much faster than the second.

It is clear that the first approach cannot be optimal. So what would you suggest to compute this kind of quantities?

Answer 1

How about

mA = RandomReal[NormalDistribution[], {3000, 3000}];
mB = RandomReal[NormalDistribution[], {3000, 3000}];
mC = RandomReal[NormalDistribution[], {3000, 3000}];

Timing[vec0 = Diagonal[mA.mC.mB];]
(*{7.515625, Null}*)

Timing[Total[Total[mA Transpose[mC.mB], {2}] - vec0]]
(*{3.968750, 3.8557602 10^-10}*)