CODEWITHSUNDEEP
pythonnumpy
Jack M
Jack M

Reputation: 5995

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

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

Upvotes: 1

Views: 1011

Answers (1)

MathIsFun7225
MathIsFun7225

Reputation: 266

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.

Upvotes: 2

Related Questions