（Python）如何获得对角线（A * B）而不必执行A * B？

Question

假设我们有两个矩阵A 和B，让矩阵C 成为A*B（矩阵乘法不是元素）。我们希望只获得C 的对角线条目，这可以通过np.diagonal(C) 完成。但是，这会导致不必要的时间开销，因为我们将 A 与 B 相乘，即使我们只需要将 A 中的每一行与具有相同“id”的 B 的列相乘，即第 1 行A 与B 的第 1 列、A 的第 2 行与B 的第 2 列等等：形成C 对角线的乘法。有没有办法使用 Numpy 有效地实现这一目标？我想避免使用循环来控制哪一行与哪一列相乘，相反，我希望有一个内置的 numpy 方法来执行这种操作以优化性能。

提前谢谢..

Answer 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 技巧有效。]

Answer 2

def diag(A,B):
    diags = []
    for x in range(len(A)):
        diags.append(A[x][x] * B[x][x])
    return diags

我相信上面的代码就是你要找的。​​p>