Numpy Add All Combinations of Vectors-CodePudding

What is the most numerically efficient way to add all the combinations of two arrays of vectors? For example what I want is the following:

a = np.array([[1,2,3], [4,5,6]])
b = np.array([[7,8,9], [10,11,12]])
[ai   bj for ai in a for bj in b]

Gives

[array([ 8, 10, 12]),
 array([11, 13, 15]),
 array([11, 13, 15]),
 array([14, 16, 18])]

It's a meshgrid with vectors instead of primary data types.

I've tried somewhat explicitly constructing the meshgrid results, which is faster than the list comprehension:

a_tile = np.tile(a, (2, 1))

array([[1, 2, 3],
       [4, 5, 6],
       [1, 2, 3],
       [4, 5, 6]])

b_repeat = np.repeat(b, 2, axis=0)

array([[ 7,  8,  9],
       [ 7,  8,  9],
       [10, 11, 12],
       [10, 11, 12]])

a_tile   b_repeat

array([[ 8, 10, 12],
       [11, 13, 15],
       [11, 13, 15],
       [14, 16, 18]])

Is this as efficient as it gets? I was looking for a way to broadcast the arrays so that the grid isn't explicitly constructed.

CodePudding user response：

You can try the following:

import numpy as np

a = np.array([[1,2,3], [4,5,6]])
b = np.array([[7,8,9], [10,11,12]])

(a[..., None]   b.T).transpose(0, 2, 1).reshape(-1, 3)

It gives:

array([[ 8, 10, 12],
       [11, 13, 15],
       [11, 13, 15],
       [14, 16, 18]])

CodePudding user response：

you can use numpy.broadcast_to to broadcast an array

N = 2 #number of repeats
your_req_array = np.broadcast_to(b.T, (N,b.shape[1], b.shape[0])).transpose(2,0,1).reshape(-1,b.shape[1])   np.broadcast_to(a, (N,a.shape[0], a.shape[1])).reshape(-1,b.shape[1])