Concatenating row-wise combinations of 2d numpy arrays into a 2d array

Question

Consider two 2d arrays, A and B.

import numpy as np

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

B = np.array([[2, 3],
              [1, 3]])

I want to build an array arrOut that gives all combinations of the rows of A and the rows of B in a 4-column array.

The desired output is:

arrOut = [[1, 4, 2, 3],
          [1, 4, 1, 3],
          [3, 5, 2, 3],
          [3, 5, 1, 3],
          [1, 2, 2, 3],
          [1, 2, 1, 3]]

I'm hoping to see a solution that could be readily expanded to all combinations of the rows of three 2d arrays to form a six column array, or all combinations of the rows of four 2d arrays to form an 8 column array.

Ehsan · Accepted Answer · 2020-08-27 21:26:57Z

Using numpy broadcasting and extendable to any number of arrays:

r1,c1 = A.shape
r2,c2 = B.shape
arrOut = np.zeros((r1,r2,c1+c2), dtype=A.dtype)
arrOut[:,:,:c1] = A[:,None,:]
arrOut[:,:,c1:] = B
arrOut.reshape(-1,c1+c2)

output:

[[1 4 2 3]
 [1 4 1 3]
 [3 5 2 3]
 [3 5 1 3]
 [1 2 2 3]
 [1 2 1 3]]

For a 3 array case (Here I used (A,B,A)):

r1,c1 = A.shape
r2,c2 = B.shape
r3,c3 = A.shape 
arrOut = np.zeros((r1,r2,r3,c1+c2+c3), dtype=A.dtype)
arrOut[:,:,:,:c1] = A[:,None,None,:]
arrOut[:,:,:,c1:c1+c2] = B[:,None,:]
arrOut[:,:,:,c1+c2:] = A
arrOut.reshape(-1,c1+c2+c3)

output:

[[1 4 2 3 1 4]
 [1 4 2 3 3 5]
 [1 4 2 3 1 2]
 [1 4 1 3 1 4]
 [1 4 1 3 3 5]
 [1 4 1 3 1 2]
 [3 5 2 3 1 4]
 [3 5 2 3 3 5]
 [3 5 2 3 1 2]
 [3 5 1 3 1 4]
 [3 5 1 3 3 5]
 [3 5 1 3 1 2]
 [1 2 2 3 1 4]
 [1 2 2 3 3 5]
 [1 2 2 3 1 2]
 [1 2 1 3 1 4]
 [1 2 1 3 3 5]
 [1 2 1 3 1 2]]

You can even make a for loop for a N array case.

Collectives™ on Stack Overflow

Concatenating row-wise combinations of 2d numpy arrays into a 2d array

1 Answer 1

Comments

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Related