在numpy包中计算矩阵乘法的对角线元素

Question

有什么方法可以计算矩阵 C 的对角线值等于 numpy.matmul(A, B) 而无需计算非对角线元素？因为我正在处理两个大矩阵 A 和 B，我只需要在我的其余代码中包含它们乘法的对角线元素。如果我计算整个矩阵 C，运行时间相当长。在此先感谢。

Answer 1

你可以这样实现：

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)

输出：

[84, 216, 366]

与使用 np.matmul

得到的对角线元素相同

print(np.matmul(arr1,arr2))

结果

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

Answer 2

你也可以使用numpy运算符einsum。阅读 the docs

对于你的特殊情况，你想使用

np.einsum('ij,ji->i', A, B)

我发现它比接受的答案更具可读性（提供对爱因斯坦符号的理解）。

此外，接受的答案returns解决方案为python list，而此答案returns为numpy array，如果您需要对此类进行操作，则这是有利的数组。

此外，此表示法允许推广到 n 维数组 (n>2)。

就运行时而言，这种方法要快得多，因为它利用 numpy 的机制来完成工作，而不是 python for 循环，后者的操作非常慢。

这两种方法自然比计算乘积然后取对角线要快得多。

测试：

小矩阵：

A = np.random.random((1000,1000))
B = np.random.random((1000,1000))

%%timeit
np.einsum('ij,ji->i', A,B)

--> 2.07 毫秒 ± 59.1 微秒

%%timeit
[sum(A[i,]*B[:,i]) for i in range(A.shape[0])]

--> 168 毫秒 ± 2.19 毫秒

更大的矩阵：

A = np.random.random((13000,13000))
B = np.random.random((13000,13000))

%%timeit
np.einsum('ij,ji->i', A,B)

--> 660 毫秒 ± 13.6 毫秒

%%timeit
[sum(A[i,]*B[:,i]) for i in range(A.shape[0])]

--> 29.5 秒 ± 516 毫秒

在numpy包中计算矩阵乘法的对角线元素

Calculation diagonal elements of matrices multiplication in numpy package

python

numpy

linear-algebra

matrix-multiplication

测试：

小矩阵：

更大的矩阵：