I have two numpy.array objects x and y where x.shape is (P, K) and y.shape is (T, K). I want to do an outer sum on these two objects such that the result has shape (P, T, K). I'm aware of the np.add.outer and the np.einsum functions but I couldn't get them to do what I wanted.

The following gives the intended result.

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

But I've got to imagine there's a faster way!

CodePudding user response：

One option is to add a new dimension to x and add using numpy broadcasting:

out = x[:, None]   y

or as @FirefoxMetzger pointed out, it's more readable to be explicit with the dimensions:

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

Test:

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