两个numpy数组的每一行组合的拼接

Question

给定两个不同形状的数组（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行，因此希望避免嵌套循环）。

- Dave Challis

Answer 1

这是一个完美的问题，可以学习切片和广播索引。

以下是使用这些工具的向量化解决方案 -

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

- Divakar

---

太棒了！ - MaxU - stand with Ukraine

谢谢，运行得很好，并且有一个额外的好处，即预先分配输出数组，因此确保我有足够的空间。 - Dave Challis

@DaveChallis 这就是整个想法。对于大型数组，任何堆叠操作都会对内存/性能的效率造成严重影响。 - Divakar

@Divakar，很漂亮的解决方案！如果您有时间，能否给我一些提示，如何使用broadcast_to函数代替强制连接的显式广播。在https://github.com/numpy/numpy/issues/2115上有一次讨论，但它基于想要连接（n,1）数组而不是整个行，并且它实际上没有使用broadcast_to。我尝试将其适应于一个数组为（3,3）而另一个为（4,2）的情况，例如[寻找形状为（12,5）的答案]。您在这里的解决方案适用于该问题，但希望更好地掌握broadcast_to()的用法。 - David R

Answer 2

参考：numpy.concatenate在记录数组中连接长度不同的字符串时失败

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)

输出：

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

- Vasili Syrakis