如何使用numpy高效计算大型二维数组各反对角线的元素和

你可以直接使用numpy的向量化接口完成计算，完全避免Python层面的双层循环，性能远高于当前的朴素实现，无需额外开发C++版本。

实现代码 #

import numpy as np

# a为输入的numpy数组
m, n = a.shape
# 生成所有元素的行列索引网格
row_idx, col_idx = np.indices((m, n))
# 计算每个元素对应的反对角线索引（同一反对角线上的元素i+j值相同）
anti_diag_idx = row_idx + col_idx
# 按反对角线索引分组求和，权重为原数组的元素值
result = np.bincount(anti_diag_idx.ravel(), weights=a.ravel())

方案说明 #

• 所有运算均为numpy底层的C实现，没有Python循环开销，针对你给出的(27584, 34092)规模的数组，只要内存充足，运行速度可以提升数千倍
• 最终输出的result形状天然为(m + n - 1,)，和需求完全匹配，无需额外调整维度
• 你可以用小规模数组验证结果正确性：比如输入a = np.array([[1,2,3],[4,5,6],[7,8,9]])，得到的结果为[ 1. 6. 15. 14. 9.]，和朴素实现的计算结果完全一致

内存紧张时的替代方案 #

如果你的设备内存不足以同时存储索引数组和原数组，可以使用逐行累加的优化方案，相比原生双层循环仍然有数量级的性能提升：

m, n = a.shape
result = np.zeros(m + n - 1, dtype=a.dtype)
for i in range(m):
    # 每行元素对应加到result的[i:i+n]位置，切片赋值运算在numpy底层实现
    result[i:i+n] += a[i]

内容的提问来源于stack exchange，提问作者user16715836