技术问询：手动实现二维数组指定维度的cumulative sum函数

一、先明确目标逻辑 #

首先得摸透numpy官方np.cumsum(arr, axis=dim)在二维数组上的行为，这是实现的基准：

• dim=0（行维度）：对每一列从上到下累加，第i行第j列的结果 = 第0到i行的第j列元素之和
• dim=1（列维度）：对每一行从左到右累加，第i行第j列的结果 = 第i行第0到j列元素之和

二、核心实现方案（禁用np.cumsum） #

方案1：循环实现（考试首选，直观易写） #

这是最适合考试场景的写法，逻辑清晰不容易出错：

1. 先获取输入数组的行列数：rows, cols = array.shape
2. 创建和输入同形状、同dtype的全零结果数组：result = np.zeros_like(array)
3. 分维度处理：
   • 当dim=0时：
     • 第一行直接复制原数组第一行：result[0] = array[0].copy()
     • 从第1行到最后一行循环：for i in range(1, rows): result[i] = result[i-1] + array[i]
   • 当dim=1时：
     • 第一列直接复制原数组第一列：result[:, 0] = array[:, 0].copy()
     • 从第1列到最后一列循环：for j in range(1, cols): result[:, j] = result[:, j-1] + array[:, j]
4. 返回result

方案2：广播/矩阵运算实现（高效进阶） #

如果想写更高效的无循环代码，可以利用下三角矩阵的特性：

• 当dim=0时：构造rows×rows的下三角全1矩阵（对角线及下方为1，上方为0），和原数组做矩阵乘法（等价于按列累加）
• 当dim=1时：构造cols×cols的下三角全1矩阵，原数组和该矩阵做乘法（等价于按行累加）
• 示例代码片段（dim=0）：

  tri_mat = np.tril(np.ones((rows, rows)))
  result = tri_mat @ array
  

三、结果验证步骤 #

必须和官方结果对比确保正确性：

1. 生成测试用例：比如test_arr = np.array([[1,2,3],[4,5,6],[7,8,9]], dtype=np.int32)
2. 分别调用自定义函数和官方函数：

   my_result_0 = cumsum_2d(test_arr, dim=0)
   np_result_0 = np.cumsum(test_arr, axis=0)
   print(np.array_equal(my_result_0, np_result_0))  # 应该输出True
   
   my_result_1 = cumsum_2d(test_arr, dim=1)
   np_result_1 = np.cumsum(test_arr, axis=1)
   print(np.array_equal(my_result_1, np_result_1))  # 应该输出True
   

四、边界情况处理 #

考试中可能会考察这些细节：

• 输入数组只有一行：dim=0时直接返回原数组
• 输入数组只有一列：dim=1时直接返回原数组
• 处理浮点型数组：确保结果dtype和输入一致，避免精度问题

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