numpy三维数组向量求和函数异常原因及修正方法

我有一个三维numpy数组：

import numpy as np
threeDimArray = np.arange(24).reshape((3, 2, 4))
print(threeDimArray)

打印结果为：

[[[ 0  1  2  3]
  [ 4  5  6  7]]
 [[ 8  9 10 11]
  [12 13 14 15]]
 [[16 17 18 19]
  [20 21 22 23]]]

定义了一个函数，想要计算数组中每个向量的和，并用该和替换原向量：

def myOperation():
    img_temp=threeDimArray.copy()
    nrows=img_temp.shape[0]
    ncolumns = img_temp.shape[1]
    for j in range(ncolumns):
        for i in range(nrows):
            img_temp[i][j]=sum(img_temp[i][j])
    return(img_temp)

预期返回：

[[6,22],
 [38,54],
 [70,86]]

实际返回：

[[[ 6,  6,  6,  6],
  [22, 22, 22, 22]],
 [[38, 38, 38, 38],
  [54, 54, 54, 54]],
 [[70, 70, 70, 70],
  [86, 86, 86, 86]]]

1. 为什么会出现这种情况？ #

• img_temp[i][j]是长度为4的一维numpy数组（对应原三维数组第i行、第j列的向量），执行img_temp[i][j] = sum(img_temp[i][j])时，numpy的广播机制会把单个求和结果扩展成和原向量长度一致的数组，再完成赋值，导致原向量的4个元素都被替换成该和。
• 代码未修改数组维度，只是更新了元素值，所以输出仍保持原三维(3,2,4)的形状，而非预期的二维(3,2)结构。

2. 如何修改函数以得到预期的返回结果？ #

有两种可行的修改方式：

方式一：创建新二维数组存储结果 #

初始化与目标形状匹配的二维数组，将每个向量的和填入对应位置：

def myOperation():
    nrows, ncolumns, _ = threeDimArray.shape
    result = np.empty((nrows, ncolumns), dtype=int)
    for j in range(ncolumns):
        for i in range(nrows):
            result[i][j] = sum(threeDimArray[i][j])
    return result

方式二：利用numpy向量化操作（更高效） #

用numpy内置求和函数指定维度，一步完成压缩计算：

def myOperation():
    return threeDimArray.sum(axis=2)

axis=2表示沿着第三个维度（即每个向量的维度）求和，直接将(3,2,4)的三维数组转换为(3,2)的二维数组，完全符合预期结果。

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