如何无需for循环为Numpy数组的子矩阵批量添加对应常量？

问题背景 #

现有一个Numpy数组X，形状为(m, n0, n1, ..., nk)，包含m个元素，每个元素都是固定形状(n0, n1, ..., nk)的矩阵；另有一个长度为m的标量数组c（例如c = np.arange(m)），要求将c[i]加到X[i]的每一个元素上。目前通过for循环实现：

for i in range(m):
    X[i] += c[i]

希望找到无需循环的向量化实现方式，类似X+c的简洁写法（但直接写X+c会因维度不匹配失败）。

解决方案 #

利用Numpy的广播机制，只需将数组c的维度扩展为与X的维度匹配，让每个c[i]能被广播到对应X[i]的所有元素上。有两种简洁实现方式：

方式1：手动扩展维度 #

针对示例中X的形状(m,4,3,6,5,7)，将c扩展为(m,1,1,1,1,1)：

X += c[:, np.newaxis, np.newaxis, np.newaxis, np.newaxis, np.newaxis]

方式2：自动适配维度（通用写法） #

如果X的维度数不确定，可通过动态生成维度的方式扩展c，无需手动写多个np.newaxis：

# 扩展c的形状为(m,1,1,...,1)，1的数量等于X除第一维外的维度数
X += c.reshape(-1, *([1]*(X.ndim - 1)))

完整示例代码 #

import numpy as np

m = 3
X = np.arange(m*4*3*6*5*7).reshape((m, 4, 3, 6, 5, 7))
c = np.arange(m)

# 替换原for循环的无循环实现
X += c.reshape(-1, *([1]*(X.ndim - 1)))

原理说明 #

直接执行X+c失败的原因是Numpy广播规则要求从最后一维开始对齐：c的形状是(m,)，X的形状是(m, n0,...)，两者最后一维长度不匹配，无法广播。而将c扩展为(m,1,1,...,1)后，其每个维度都能与X对应维度对齐（1可以匹配任意长度的维度），因此每个c[i]会被自动广播到X[i]的所有元素上，实现批量加法。

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