我可以einsum在这里使用：
>>> 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])
对于较大的数组，它比直接进行乘法要快得多：
>>> 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
[注意：起初我是做元素级版本ii,ii->i，而不是矩阵乘法。相同的einsum技巧起作用。]
602392714
清零编程群