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

user_815612 0
Viewed 149

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

numpy python
1 Answers

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.

user_115134330
Related