如何对3D NumPy数组提取2D对角线求和并得到shape为(3,10)的结果

错误原因 #

你原有代码遍历的偏移量范围是range(XRN_LEN)，也就是range(4)，对应你选择的两个对角线轴（维度分别为3和4）的4个正偏移，所以输出形状是(4,10)，和预期不符。

实现代码 #

如果你要保留原有跨第一、第二维度取对角线的逻辑，仅调整输出形状为(3,10) #

只要把偏移量遍历范围改成第一维的长度3即可：

RIGHT = [XRN.diagonal(i, axis1=0, axis2=1).sum(axis=1) for i in range(XRN.shape[0])]
np_RIGHT = np.array(RIGHT)
print(np_RIGHT.shape) # 输出 (3, 10)

如果你是要对每个独立的2D切片（每个第一维对应的4行10列矩阵）分别求对角线和 #

如果你需要的是3个2D矩阵各自的10个对角线偏移的求和结果，用以下代码：

result = []
for single_mat in XRN:
    # 遍历10个偏移量，覆盖所有列方向的对角线
    diag_sum_list = [single_mat.diagonal(offset=o).sum() for o in range(single_mat.shape[1])]
    result.append(diag_sum_list)
result = np.array(result)
print(result.shape) # 输出 (3, 10)

补充 #

如果需要计算反对角线和，可先将矩阵左右翻转再取对角线：np.fliplr(single_mat).diagonal(offset=o).sum()

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