自定义NumPy数组求和函数报错排查：1D/2D/3D数组求和异常

错误原因分析 #

原函数的核心问题是没有正确遍历数组的所有元素，而是每次循环都累加同一个固定位置的元素：

• 1D数组：循环次数等于数组长度，但每次累加array[0]（第一个元素），相当于把第一个元素重复加了N次，而非遍历每个元素。
• 2D数组：仅循环列数，每次累加array[0,1]（第一行第二列的固定元素），完全没有遍历所有行和列。
• 3D数组：仅循环最内层维度的长度，每次累加array[0,1,2]（固定位置元素），没有遍历所有块、行、列。
  另外，函数内使用sum作为变量名会覆盖Python内置的sum函数，虽不影响运行，但属于不良编码习惯。

修正后的代码 #

import numpy as np

array_1d = np.array([1,2,3,4,5])
array_2d = np.array([[1,3,4,5],[1,4,5,7]])
array_3d = np.array([[[1,3,4,5],[1,4,5,7]],[[6,7,8,9],[7,8,7,6]]])

def custom_sum(arr):
    total = 0
    shape = arr.shape
    dim = len(shape)
    if dim == 3:
        # 遍历3D数组的块、行、列
        for i in range(shape[0]):
            for j in range(shape[1]):
                for k in range(shape[2]):
                    total += arr[i,j,k]
    elif dim == 2:
        # 遍历2D数组的行、列
        for i in range(shape[0]):
            for j in range(shape[1]):
                total += arr[i,j]
    elif dim == 1:
        # 遍历1D数组的每个元素
        for i in range(shape[0]):
            total += arr[i]
    return total

# 测试验证
print(custom_sum(array_1d))  # 输出15
print(custom_sum(array_2d))  # 输出30
print(custom_sum(array_3d))  # 输出88

通用优化版本（支持任意维度） #

如果想让函数适配任意维度的数组，可使用递归方式逐层拆解求和，避免针对每个维度写嵌套循环：

def recursive_sum(arr):
    total = 0
    if arr.ndim == 1:
        for num in arr:
            total += num
    else:
        for sub_arr in arr:
            total += recursive_sum(sub_arr)
    return total

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