python - 沿指定轴的两个 numpy 数组的外部总和

Question

我有两个 numpy.array 对象 x 和 y 其中 x.shape 是 ( P, K) 和 y.shape 是 (T, K)。我想对这两个对象进行外部求和，使结果具有 (P, T, K) 的形状。我知道 np.add.outer 和 np.einsum 函数，但我无法让它们执行我想要的操作。

下面给出了预期的结果。

  x_plus_y = np.zeros((P, T, K))
  for k in range(K):
    x_plus_y[:, :, k] = np.add.outer(x[:, k], y[:, k])

但我必须想象有更快的方法!

Answer 1

一个选项是向 x 添加一个新维度并使用 numpy 广播添加:

out = x[:, None] + y

或作为 @FirefoxMetzger pointed out ，更明确的维度更具可读性:

out = x[:, None, :] + y[None, :, :]

测试:

P, K, T = np.random.randint(10,30, size=3)
x = np.random.rand(P, K)
y = np.random.rand(T, K)
x_plus_y = np.zeros((P, T, K))
for k in range(K):
    x_plus_y[:, :, k] = np.add.outer(x[:, k], y[:, k])

assert (x_plus_y == x[:, None] + y).all()