How can I form the matrix of sums of two arrays of vectors in numpy?

Question

I have two numpy arrays a and b , of lengths n and m , containing vectors of length d , that is, a[i].shape == (d,) and likewise for b , for all i .

I want to form the matrix A of sums of these vectors, so that

A[i, j] == a[i] + b[j]

so A will be of shape (n, m, d) .

I know that I can use numpy.add.outer to get a matrix M of shape (N, d, M, d) containing all possible sums of all components of all vectors in a and b , and then my matrix A would be given by the formula

A[i, j, k] = M[i, k, j, k]

However, other than writing for loops I don't know how to turn M into what I need, and I need to do this in the most efficient way possible.

Ultimately I want to end up with the matrix of sums of squares of these vectors by doing (A**2).sum(axis=2) .

Answer 1

You can utilize broadcasting :

A = a[:, None, :] + b[None, :, :]

You could also replace None with numpy.newaxis . They are the same thing, but numpy.newaxis might be easier to understand. I will note that you can also just do

A = a[:, None, :] + b

but I think the other way makes it more clear what is going on.