python - 两个 numpy 数组中所有行的组合

Question

我有两个数组，例如 (3,2) 的形状和另一个 (10,7) 的形状。我想要两个数组的所有组合，这样我最终得到一个 9 列数组。换句话说，我想要第一个数组的每一行与第二个数组的行的所有组合。

我该怎么做？据我所知，我没有正确使用 meshgrid。

根据之前的帖子，我的印象是

a1 = np.zeros((10,7))
a2 = np.zeros((3,2))
r = np.array(np.meshgrid(a1, a2)).T.reshape(-1, a1.shape[1] + a2.shape[1])

会起作用，但这给了我 (84,10) 的维度。

Answer 1

方法#1

着重于性能，这是一种使用 array-initialization 和 element-broadcasting 进行分配的方法 -

m1,n1 = a1.shape
m2,n2 = a2.shape
out = np.zeros((m1,m2,n1+n2),dtype=int)
out[:,:,:n1] = a1[:,None,:]
out[:,:,n1:] = a2
out.shape = (m1*m2,-1)

解释:

诀窍在于两个步骤:

out[:,:,:n1] = a1[:,None,:]
out[:,:,n1:] = a2

第 1 步:

In [227]: np.random.seed(0)

In [228]: a1 = np.random.randint(1,9,(3,2))

In [229]: a2 = np.random.randint(1,9,(2,7))

In [230]: m1,n1 = a1.shape
     ...: m2,n2 = a2.shape
     ...: out = np.zeros((m1,m2,n1+n2),dtype=int)
     ...: 

In [231]: out[:,:,:n1] = a1[:,None,:]

In [232]: out[:,:,:n1]
Out[232]: 
array([[[5, 8],
        [5, 8]],

       [[6, 1],
        [6, 1]],

       [[4, 4],
        [4, 4]]])

In [233]: a1[:,None,:]
Out[233]: 
array([[[5, 8]],

       [[6, 1]],

       [[4, 4]]])

所以，基本上我们是在分配 a1 的元素，保持第一个轴与相应的输出轴对齐，同时让输出数组的第二个轴上的元素填充到广播中与沿该轴为 a1 添加的 newaxis 相对应的方式。这是这里的关键，它带来了性能，因为我们没有分配额外的内存空间，否则我们需要使用显式重复/平铺方法。

第 2 步:

In [237]: out[:,:,n1:] = a2

In [238]: out[:,:,n1:]
Out[238]: 
array([[[4, 8, 2, 4, 6, 3, 5],
        [8, 7, 1, 1, 5, 3, 2]],

       [[4, 8, 2, 4, 6, 3, 5],
        [8, 7, 1, 1, 5, 3, 2]],

       [[4, 8, 2, 4, 6, 3, 5],
        [8, 7, 1, 1, 5, 3, 2]]])

In [239]: a2
Out[239]: 
array([[4, 8, 2, 4, 6, 3, 5],
       [8, 7, 1, 1, 5, 3, 2]])

在这里，我们基本上是沿输出数组的第一个轴广播那个 block a2，而没有明确地进行重复复制。

示例输入，完整性输出 -

In [242]: a1
Out[242]: 
array([[5, 8],
       [6, 1],
       [4, 4]])

In [243]: a2
Out[243]: 
array([[4, 8, 2, 4, 6, 3, 5],
       [8, 7, 1, 1, 5, 3, 2]])

In [244]: out
Out[244]: 
array([[[5, 8, 4, 8, 2, 4, 6, 3, 5],
        [5, 8, 8, 7, 1, 1, 5, 3, 2]],

       [[6, 1, 4, 8, 2, 4, 6, 3, 5],
        [6, 1, 8, 7, 1, 1, 5, 3, 2]],

       [[4, 4, 4, 8, 2, 4, 6, 3, 5],
        [4, 4, 8, 7, 1, 1, 5, 3, 2]]])

方法#2

另一个平铺/重复 -

parte1 = np.repeat(a1[:,None,:],m2,axis=0).reshape(-1,m2)
parte2 = np.repeat(a2[None],m1,axis=0).reshape(-1,n2)
out = np.c_[parte1, parte2]