How can I extract only the diagonal elements of a matrix product without computing the whole product matrix?

I have nxk matrix (A) and a kxk (B) matrix where n>>k. I want the diagonal terms of A*B*A' in a n-vector. It would be wasteful to compute the whole A*B*A' and then extract the diagonal vector. What is the easy and efficient way to do this?

Thanks.

[Cedric]

How large are n and k?

[Ranjan Sonalkar]

n can be 100000 or 1000000. k is 10 or less.

Answer 1

[David Young]

You could use

sum((a * b) .* a, 2)

A check:

a = rand(1000, 10);
b = rand(10);
max(abs(diag(a*b*a.') - sum((a*b).*a,2)))

And yes, it's quicker:

f1 = @() diag(a * b * a.');
f2 = @() sum((a*b).*a, 2);

timeit(f1)

ans = 0.0104

timeit(f2)

ans = 1.2795e-04

[Edit: my example originally had n < k, which only gave a modest speedup. For n >> k, as required, the speedup is much greater, as expected.]

[VICTOR MIRASIERRA CALLEJA]

This helped a lot. Thank you!