如何在NumPy中按元素将一维数组与二维数组相加得到三维数组

当然可以用纯NumPy实现，核心是利用NumPy的广播机制，无需编写循环，效率更高。

实现代码 #

import numpy as np

A = np.array([
    [10, 9, 8, 7, 6],
    [5, 4, 3, 2, 1]
])
B = np.array([1, 2, 3])

# 调整B的形状以触发广播，与A相加得到目标三维数组
result = A + B.reshape(-1, 1, 1)

# 验证结果
print(result)

输出结果 #

[[[11 10  9  8  7]
  [ 6  5  4  3  2]]

 [[12 11 10  9  8]
  [ 7  6  5  4  3]]

 [[13 12 11 10  9]
  [ 8  7  6  5  4]]]

原理说明 #

NumPy的广播机制会自动扩展数组形状，使维度匹配的数组可以进行元素级运算：

• A的形状为 (2, 5)（2行5列的二维数组）
• 通过B.reshape(-1, 1, 1)将B的形状调整为 (3, 1, 1)（三维数组，3个1x1的元素）
• 相加时，B的后两个维度会自动扩展为(2,5)，同时A会被扩展第一个维度为3，最终实现每个B的元素与整个A相加，生成(3,2,5)的三维数组。

替代写法 #

也可以用np.expand_dims来调整B的形状，效果完全一致：

result = A + np.expand_dims(B, axis=(1, 2))

这种方法完全依赖NumPy的底层优化，比Python循环的计算效率高得多，尤其适合处理大规模数组。

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