在numpy中计算对角线和[重复]

Posted 2023-03-12

【中文标题】在numpy中计算对角线和[重复]

【英文标题】：Computing diagonal sums in numpy [duplicate]

【发布时间】：2022-01-08 16:26:28

【问题描述】：

我有一个相当大的矩形 numpy 数组，形状为 (m, n)，例如：

>>> a.shape
(27584, 34092)

我必须计算数组中每个 anti-diagonal 的总和。这个新数组的形状为(m + n - 1,)。

简单的方法是这样做：

m, n = a.shape
r = np.zeros(m + n - 1)
for i in range(m):
    for j in range(n):
        r[i + j] += a[i][j]

# r is the sum of all anti-diagonals of a

这显然很慢，有没有办法使用聪明的 numpy 原语来执行计算？我唯一的其他选择是用 C++ 编写代码，这也是可行的 - 但需要更多的工作。

【参考方案1】：

在我看来，最直观的解决方案是：

使用 a[::-1] 来“翻转”行的顺序，这样你就可以 在“正常”对角线上操作， 使用 np.diagonal 访问每个对角线， 对返回的对角线求和， 将上述代码放入列表推导式中。

所以整个代码可以是：

result = [np.diagonal(a[::-1], offs).sum() for offs in range(-a.shape[0] + 1, a.shape[1])]

它的运行速度应该比您的代码快得多。

由于您的源数组相当大，也许您会获得一些速度 获得：

在开始时计算“翻转”数组

一次

， 然后对这个数组进行操作：

所以重新编写的代码，包括关于 np.trace 的提示（参见 cmets）， 可以是：

b = a[::-1]
result = [np.trace(b, offs) for offs in range(-a.shape[0] + 1, a.shape[1])]

【讨论】：

您可以使用np.trace() 作为np.diagonal(...).sum() 的快捷方式。

我最终在 C++ 中重做了它，以确保它对于更大的矩阵来说是快速的并且代码不会太复杂。但是，这个解决方案无论如何都很好，所以我接受了它。 @Valdi_bo