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

Question

我有两个numpy.array对象x，y其中x.shapeis(P, K)和y.shapeis (T, K)。我想对这两个对象进行外部求和，以使结果具有 shape (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 指出的那样，明确的尺寸更具可读性：

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()