如何基于其他维度索引界定的范围对ndarray进行求和？

针对你提出的需求——沿部分维度求和，且求和范围由其他维度的索引定义，这里提供一种基于前缀和+向量化索引的高效实现方案，完全避免Python层面的循环，适合大规模数组场景。

实现思路 #

• 前缀和预处理：先对目标维度（这里是第三维，axis=2）计算前缀和，任意连续区间的和可通过前缀和的差值快速得到。
• 生成索引网格：构造第一维（i）和第二维（j）的索引网格，对应每个需要计算的(i,j)对。
• 向量化计算区间和：利用索引网格直接从前缀和数组中取出对应位置的值，计算差值，并处理i≥j时的空切片情况（和为0）。

代码实现 #

import numpy as np

# 构造示例数组
x = np.arange(2*3*4).reshape((2,3,4))

# 1. 计算第三维的前缀和（前面补0，方便计算区间[0,k)的和）
pre_sum = np.concatenate(
    [np.zeros((x.shape[0], x.shape[1], 1)), np.cumsum(x, axis=2)],
    axis=2
)

# 2. 生成i和j的索引网格（indexing='ij'保证网格形状与原数组前两维匹配）
i_grid, j_grid = np.meshgrid(np.arange(x.shape[0]), np.arange(x.shape[1]), indexing='ij')

# 3. 计算每个(i,j)对应的区间和，i≥j时置为0
result = np.where(
    i_grid < j_grid,
    pre_sum[i_grid, j_grid, j_grid] - pre_sum[i_grid, j_grid, i_grid],
    0
)

print(result)
# 输出：
# [[ 0  4 17]
#  [ 0  0 21]]

方案优势 #

• 效率极高：所有核心操作（cumsum、索引、where）都是numpy的向量化内置函数，基于C实现，比Python循环/列表推导快几个数量级，尤其适合大规模数组。
• 扩展性强：如果需求变化（比如求和维度改变、索引规则调整），只需修改前缀和的axis参数或索引网格的生成逻辑即可。

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