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

Question

我有一个相当大的矩形 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++ 编写代码，这也是可行的 - 但需要更多的工作。

Answer 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])]