Numpy数组不同长度轴求和的高效优雅方案问询

问题背景 #

给定形状为 r×c×d 的Numpy 3D张量 D，以及形状为 r×c 的整数矩阵 Q（所有元素小于d），需要对每个 D[r_i][c_i] 的第三维前 Q[r_i][c_i]+1 个元素求和，得到 r×c 的结果矩阵。

示例数据：

import numpy as np

r = 2
c = 3
d = 3    
D = np.array([[[1, 5, 3], [1, 2, 5], [1, 4, 3]], [[1, 1, 6], [3, 1, 7], [5, 1, 3]]])
Q = np.array([[1, 1, 2], [2, 1, 2]])

预期输出：

Res = np.array([[6, 3, 8], [8, 4, 5]])

现有实现为嵌套列表推导：

res = np.array([[np.sum(D[r_i, c_i, :Q[r_i, c_i]+1]) for c_i in range(c)] for r_i in range(r)])

更高效优雅的解决方案 #

方法1：前缀和数组 + 索引提取 #

这是效率最高的方案，完全基于Numpy矢量化操作，避免Python循环：

# 计算第三维的前缀和，形状保持(r, c, d)
prefix_sum = D.cumsum(axis=2)
# 构造广播索引，提取Q对应位置的前缀和结果
res = prefix_sum[np.arange(r)[:, None], np.arange(c), Q]

原理：

• D.cumsum(axis=2) 会对每个D[r_i][c_i]的第三维元素累加，比如D[0][0] = [1,5,3]的前缀和为[1,6,9]，Q[0][0] = 1对应的就是前2个元素的和6，正好匹配需求。
• np.arange(r)[:, None] 和 np.arange(c) 用于构造广播索引，让Q的每个元素精准对应到prefix_sum的目标位置。

方法2：掩码矩阵求和 #

通过构造掩码标记需要求和的元素，再整体求和，可读性较好：

# 生成第三维索引矩阵，转置后与D形状匹配
idx = np.arange(d)[None, None, :]
# 构造掩码：第三维索引 <= Q对应位置的元素
mask = idx <= Q[:, :, None]
# 掩码过滤后沿第三维求和
res = (D * mask).sum(axis=2)

原理：

• 先生成和D同形状的布尔掩码矩阵，True表示该元素需要计入求和，False则忽略。
• 用D * mask将不需要的元素置0，最后沿第三维求和得到每个位置的结果。

方案对比 #

• 列表推导方案：本质是Python循环，当r、c较大时，频繁调用np.sum会产生大量Python与Numpy的交互开销，效率显著下降。
• 前缀和方案：纯Numpy内部矢量化操作，速度最快，尤其适合大张量场景。
• 掩码方案：同样是矢量化操作，但需要额外存储掩码矩阵，内存占用略高，速度略逊于前缀和方案，但代码逻辑更直观。

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