Question

In NumPy:

A = np.array([[1,2,3],[4,5,6]])
array([[1, 3, 5],
       [2, 4, 6]])

B = np.array([[1,2],[3,4],[5,6]])
array([[1, 2],
       [3, 4],
       [5, 6]])

A.dot(B)
array([[35, 44],
       [44, 56]])

I only care about getting A.dot(B).diagonal() = array([35, 56])

Is there a way I can get array([35, 56]) without having to compute the inner products of all the rows and columns? I.e. the inner product of the ith row with ith column?

I ask because the performance difference becomes more significant for larger matrices.

Answer 1

This is just matrix multiplication for 2D arrays:

C[i, j] = sum(A[i, ] * B[, j])

So since you just want the diagonal elements, looks like you’re after

sum(A[i, ] * B[, i]) # for each i

So you could just use list comprehension:

[np.dot(A[i,:], B[:, i]) for i in xrange(A.shape[0])]
# [22, 64]

OR, (and this only works because you want a diagonal so this assumes that if A’s dimensions are n x m, B’s dimensions will be m x n):

np.sum(A * B.T, axis=1)
# array([22, 64])

(no fancy numpy tricks going on here, just playing around with the maths).