我可能会在这里使用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
技巧也可以使用。]
提示:原来我用的逐元素版本是指对于两个形状相等的数组,对它们的每个相应元素进行操作,例如将它们相加或相乘。而矩阵乘法则涉及到两个矩阵相乘生成一个新的矩阵的操作。
A*B
表示逐元素相乘,而不是矩阵乘法(矩阵乘法应该使用a.dot(b)
)。 - BlairA
和B
是numpy.array
,那就是这种情况。如果它们是numpy.matrix
,你可以使用A*B
。 - John La Rooynumpy.matrix
的存在。感谢你指出这一点。 - BlairA
和B
是矩阵,抱歉没有说明清楚。 - IssamLaradji