小编典典

(Python)如何在不执行A * B的情况下获得对角线(A * B)?

python

假设我们有两个矩阵A和,B而矩阵CA*B(矩阵乘法而不是元素方式)。我们希望仅获取的对角线条目C,可以通过来完成np.diagonal(C)。然而,这会导致不必要的时间开销,因为我们正在为B乘以即使我们只需要每一行的乘法在A与列B具有相同的“身份证”,即第1行A与第1列B,行2
ofA与2的列等B,依此类推:形成对角线的乘法C。有没有办法使用Numpy有效地实现这一目标?我想避免使用循环来控制哪一行乘以哪一列,相反,我希望使用内置的numpy方法来执行这种操作以优化性能。

提前致谢..


阅读 120

收藏
2021-01-20

共1个答案

小编典典

我可以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技巧起作用。]

2021-01-20