CODEWITHSUNDEEP
pythonarraysnumpyvectorization
vedran
vedran

Reputation: 781

Combine array along axis

I'm trying to do the following:

I have a (4,2)-shaped array:

a = np.array([[-1, 0],[1, 0],[0, -1], [0, 1]])

I have another (2, 2)-shaped array:

b = np.array([[10, 10], [5, 5]])

I'd like to add them along rows of b and concatenate, so that I end up with:

[[ 9, 10], 
 [11, 10],
 [10,  9],
 [10, 11],
 [4, 5],
 [6, 5],
 [5, 4],
 [5, 6]]

The first 4 elements are b[0]+a, and the last four are b[1]+a. How can i generalize that if b is (N, 2)-shaped, not using a for loop over its elements?

Upvotes: 2

Views: 52

Answers (1)

Divakar
Divakar

Reputation: 221584

You can use broadcasting to get all the summations in a vectorized manner to have a 3D array, which could then be stacked into a 2D array with np.vstack for the desired output. Thus, the implementation would be something like this -

np.vstack((a + b[:,None,:]))

Sample run -

In [74]: a

Out[74]: 
array([[-1,  0],
       [ 1,  0],
       [ 0, -1],
       [ 0,  1]])

In [75]: b
Out[75]: 
array([[10, 10],
       [ 5,  5]])

In [76]: np.vstack((a + b[:,None,:]))
Out[76]: 
array([[ 9, 10],
       [11, 10],
       [10,  9],
       [10, 11],
       [ 4,  5],
       [ 6,  5],
       [ 5,  4],
       [ 5,  6]])

You can replace np.dstack with some reshaping and this might be a bit more efficient, like so -

(a + b[:,None,:]).reshape(-1,a.shape[1])

Upvotes: 3

Related Questions