(Python) How to get diagonal(A*B) without having to perform A*B?

Question

Let\'s say we have two matrices A and B and let matrix C be A*B (matrix multiplication not element-wise). We wish to get

Answer 1

def diag(A,B):
    diags = []
    for x in range(len(A)):
        diags.append(A[x][x] * B[x][x])
    return diags

I believe the above code is that you're looking for.

Answer 2

I might use einsum here:

>>> a = np.random.randint(0, 10, (3,3))
>>> b = np.random.randint(0, 10, (3,3))
>>> a
array([[9, 2, 8],
       [5, 4, 0],
       [8, 0, 6]])
>>> b
array([[5, 5, 0],
       [3, 5, 5],
       [9, 4, 3]])
>>> a.dot(b)
array([[123,  87,  34],
       [ 37,  45,  20],
       [ 94,  64,  18]])
>>> np.diagonal(a.dot(b))
array([123,  45,  18])
>>> np.einsum('ij,ji->i', a,b)
array([123,  45,  18])

For larger arrays, it'll be much faster than doing the multiplication directly:

>>> a = np.random.randint(0, 10, (1000,1000))
>>> b = np.random.randint(0, 10, (1000,1000))
>>> %timeit np.diagonal(a.dot(b))
1 loops, best of 3: 7.04 s per loop
>>> %timeit np.einsum('ij,ji->i', a, b)
100 loops, best of 3: 7.49 ms per loop

[Note: originally I'd done the elementwise version, ii,ii->i, instead of matrix multiplication. The same einsum tricks work.]