两个numpy数组的每个行组合的连接

给定两个不同形状的数组(A和B),我想生成一个数组,其中包含A中每一行与来自B的每一行的串联.

例如给出:

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

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

想生成数组:

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

我可以通过迭代来做到这一点,但它非常慢,所以寻找numpy能够尽可能高效地重新创建上述函数的一些函数组合(输入数组A和B的大小最多为10,000行,因此希望避免嵌套循环).

学习slicing和完善的完美问题broadcasted-indexing.

这是使用这些工具的矢量化解决方案 -

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)

样品运行 -

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