[FIXED] Calculation diagonal elements of matrices multiplication in numpy package

Issue

Is there any way to calculate the diagonal value of matrix C which is equal to numpy.matmul(A, B) without the need to calculation off-diagonal elements? Because I am working with two large matrices A and B and I only need to have diagonal elements of their multiplication in the rest of my code. And the run time is considerably large if I calculate the whole matrix C. Thanks in advance.

Solution

You can achieve it as :

arr1 = np.array([[1,2,3],[4,5,6],[7,8,9]])
arr2 = np.array([[10,11,12],[13,14,15],[16,17,18]])

diag_elem = [sum(arr1[i,]*arr2[:,i]) for i in range(arr1.shape[0])]
print(diag_elem)

output:

[84, 216, 366]

which are the same diagonal elements you get using the np.matmul

print(np.matmul(arr1,arr2))

results in

array([[ 84,  90,  96],
   [201, 216, 231],
   [318, 342, 366]])

Answered By - Arvind Kumar