Concatenation of every row combination of two numpy arrays

Question

Given two arrays (A and B) of different shapes, I've like to produce an array containing the concatenation of every row from A with every row from B.

E.g. given:

A = np.array([[1, 2],
              [3, 4],
              [5, 6]])

B = np.array([[7, 8, 9],
              [10, 11, 12]])

would like to produce the array:

[[1, 2, 7, 8, 9],
 [1, 2, 10, 11, 12],
 [3, 4, 7, 8, 9],
 [3, 4, 10, 11, 12],
 [5, 6, 7, 8, 9],
 [5, 6, 10, 11, 12]]

I can do this with iteration, but it's very slow, so looking for some combination of numpy functions that can recreate the above as efficiently as possible (the input arrays A and B will be up to 10,000 rows in size, hence looking to avoid nested loops).

Answer 1

Perfect problem to learn about slicing and broadcasted-indexing.

Here's a vectorized solution using those tools -

def concatenate_per_row(A, B):
    m1,n1 = A.shape
    m2,n2 = B.shape

    out = np.zeros((m1,m2,n1+n2),dtype=A.dtype)
    out[:,:,:n1] = A[:,None,:]
    out[:,:,n1:] = B
    return out.reshape(m1*m2,-1)

Sample run -

In [441]: A
Out[441]: 
array([[1, 2],
       [3, 4],
       [5, 6]])

In [442]: B
Out[442]: 
array([[ 7,  8,  9],
       [10, 11, 12]])

In [443]: concatenate_per_row(A, B)
Out[443]: 
array([[ 1,  2,  7,  8,  9],
       [ 1,  2, 10, 11, 12],
       [ 3,  4,  7,  8,  9],
       [ 3,  4, 10, 11, 12],
       [ 5,  6,  7,  8,  9],
       [ 5,  6, 10, 11, 12]])