CODEWITHSUNDEEP
pythonperformancenumpyvectorization
Dave Challis
Dave Challis

Reputation: 3906

Concatenation of every row combination of two numpy arrays

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

Upvotes: 4

Views: 1846

Answers (2)

Divakar
Divakar

Reputation: 221534

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

Upvotes: 7

Vasili Syrakis
Vasili Syrakis

Reputation: 9591

Reference: numpy.concatenate on record arrays fails when array has different length strings

import numpy as np
from numpy.lib.recfunctions import stack_arrays
from pprint import pprint

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

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

cartesian = [stack_arrays((a, b), usemask=False) for a in A
                                                 for b in B]

pprint(cartesian)

Output:

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

Upvotes: 2

Related Questions