Python中二维数组按列从指定起始位置求和的实现方案

先把你给出的示例数据再明确一下，方便后续验证：

import numpy as np
a = np.arange(20).reshape(5, 4)
ref = np.array([0, 2, 4, 1])

运行后a的内容是：

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

ref的每个元素对应每一列的求和起始行索引，比如第1列（索引1）要从行2开始求和，第3列（索引3）从行1开始求和。

方法一：矢量化掩码法（适合大型数组） #

这种方法利用numpy的广播机制，全程矢量化操作，效率很高，适合处理大规模数据：

# 生成和a同形状的行索引矩阵
row_indices = np.arange(a.shape[0])[:, np.newaxis]
# 创建掩码：标记哪些行需要参与对应列的求和
mask = row_indices >= ref
# 把不需要的元素置0后按列求和
result = np.where(mask, a, 0).sum(axis=0)

运行后result的输出是：array([40, 39, 18, 52])，完全符合预期：

• 第0列（起始行0）：0+4+8+12+16 = 40
• 第1列（起始行2）：9+13+17 = 39
• 第2列（起始行4）：18 = 18
• 第3列（起始行1）：7+11+15+19 = 52

方法二：列表推导式（直观易懂） #

如果你的数组规模不大，用列表推导式会更直观，代码可读性更强：

result = np.array([a[ref[i]:, i].sum() for i in range(a.shape[1])])

这个逻辑很简单：遍历每一列的索引i，对a的第i列从ref[i]行开始切片，然后求和，最后把所有列的结果拼成数组，输出和上面完全一致。

两种方法各有优势，你可以根据自己的数据规模和代码可读性需求来选择~

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