Trace of Matrix Product

Problem

Compute the trace of the matrix product A*B without forming the full product:

tr(A*B) = sum_i sum_k A[i][k] * B[k][i]

A is n x m and B is m x n, so the product is square.

Hint

Only diagonal entries of the product matter, so skip the rest.

case 1

in: [[[1,2],[3,4]],[[5,6],[7,8]]]

expected: 69

case 2

in: [[[1,2,3],[4,5,6]],[[1,0],[0,1],[1,1]]]

expected: 15

case 3

in: [[[1,0],[0,1]],[[7,8],[9,10]]]

expected: 17

case 4

in: [[[0,0],[0,0]],[[5,6],[7,8]]]

expected: 0

case 5

in: [[[2.5,1],[0.5,2]],[[1,0.5],[0.25,2]]]

expected: 7

Starter codePython

def trace_of_product(A, B):
    # Your code here
    pass

5 test cases · run your code against them in the editor