（Python）如何在不执行A*B的情况下获得对角线（A*B）？

Question

假设我们有两个矩阵A和B，并让矩阵C为A*B（矩阵乘法而不是元素逐个相乘）。我们希望仅获取C的对角线条目，可以通过np.diagonal(C)来实现。然而，这会导致不必要的时间开销，因为我们正在将A与B相乘，尽管我们只需要A中每一行与具有相同'id'的B列相乘，即A的第1行与B的第1列相乘，A的第2行与B的第2列相乘，以此类推：组成C对角线的乘积。有没有一种使用Numpy高效实现这样操作的方法？我想避免使用循环来控制哪行与哪列相乘，而希望使用内置的numpy方法来优化性能。

提前感谢。

- IssamLaradji

---

对于任何查看此内容的人，需要注意：在NumPy中，A*B表示逐元素相乘，而不是矩阵乘法（矩阵乘法应该使用a.dot(b)）。 - Blair

"A"和"B"是"ndarray"还是"matrix"类型？ - Bitwise

@Blair，如果A和B是numpy.array，那就是这种情况。如果它们是numpy.matrix，你可以使用A*B。 - John La Rooy

@gnibbler 哎呀，我已经习惯于处理三维数据，以至于忘记了 numpy.matrix 的存在。感谢你指出这一点。 - Blair

@Bitwise，A和B是矩阵，抱歉没有说明清楚。 - IssamLaradji

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 技巧也可以使用。]

提示：原来我用的逐元素版本是指对于两个形状相等的数组，对它们的每个相应元素进行操作，例如将它们相加或相乘。而矩阵乘法则涉及到两个矩阵相乘生成一个新的矩阵的操作。

- DSM

Answer 2

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

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

- BenjaminCohen

---

这不是矩阵乘法的工作方式。 - John La Rooy

抱歉，不知道你是指标量积还是矩阵积。 - BenjaminCohen