Numpy广播实现多维数组减法时输出形状不符的原因咨询

问题代码 #

import numpy as np
data = np.array([[1, 2], [3, 4]])
a = np.array([[0, 0], [1, 1]])
p = data - a[:, np.newaxis]
print(p)

实际输出 #

[[[1 2],
  [3 4]],
 [[0 1],
  [2 3]]]

预期输出 #

[[[1,2]],
 [[2,3]]]

为什么会得到当前输出？ #

核心是对Numpy广播规则的理解偏差，拆解形状和广播过程就能明白：

1. 各数组形状
   • data：形状为 (2, 2)（2行2列）
   • a[:, np.newaxis]：给a新增一个中间维度，形状变为 (2, 1, 2)（包含2个1行2列的子数组）
2. 广播匹配逻辑
   Numpy广播从最后一维开始对齐，维度长度不匹配时会自动扩展长度为1的维度：
   • data 的形状 (2,2) 会被隐式升维为 (1, 2, 2)
   • a[:, np.newaxis] 的形状是 (2,1,2)
   • 两者广播后都会被扩展为 (2, 2, 2)：
     • data 复制第一维度（从1→2），相当于把整个data重复2次
     • a[:, np.newaxis] 复制中间维度（从1→2），相当于把每个子数组重复2次
   • 逐元素相减后，得到形状为 (2,2,2) 的数组，也就是你看到的输出。

如何得到预期输出？ #

如果想让a的每一行和data的对应行相减，再整理成(2,1,2)的形状，直接对原数组做对应行减法后新增维度即可：

import numpy as np
data = np.array([[1, 2], [3, 4]])
a = np.array([[0, 0], [1, 1]])
p = (data - a)[:, np.newaxis]
print(p)

输出结果：

[[[1 2]]
 [[2 3]]]

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