NumPy不同维度数组相加如何避免重复书写np.newaxis

你不需要手动重复书写多次np.newaxis，以下几种写法可以自动适配任意维度的数组A，无需手动计数需要补充的新轴数量： #

方法1：reshape自动补全后置单维度 #

直接计算两个数组的维度差，自动生成对应数量的长度为1的后置维度，适配任意维度的A：

import numpy as np
A = np.zeros([20,2,2,2,2,3])
B = np.arange(1,21)
# B调整后形状为(20, 1, 1, ..., 1)，单维度数量和A除第一维外的维度数一致
B = B.reshape(B.shape[0], *([1] * (A.ndim - B.ndim)))
C = A + B

方法2：np.expand_dims批量插入新轴 #

np.expand_dims支持传入轴序号元组，一次性在所有需要的位置插入新轴，逻辑更直观：

# 在B的第1位到最后一位，批量插入长度为1的新轴
B = np.expand_dims(B, axis=tuple(range(1, A.ndim)))
C = A + B

方法3：批量生成np.newaxis索引 #

如果你习惯用np.newaxis（别名None）的索引写法，可以通过元组乘法一次性生成所有需要的新轴索引，避免重复书写：

# 第一个索引全选B的原有维度，后面拼接对应数量的np.newaxis
B = B[(slice(None),) + (np.newaxis,) * (A.ndim - B.ndim)]
C = A + B

补充说明：不需要把B真正扩展为和A完全一致的形状（比如复制为(20,2,2,2,2,3)），只要调整为广播兼容的形状(20,1,1,1,1,1)即可。NumPy运算时会自动对长度为1的维度做广播处理，不会额外复制数据占用内存，运行效率更高。

内容的提问来源于stack exchange，提问作者Ku-trala