连接两个 numpy 数组的每一行组合答案

【问题标题】：Concatenation of every row combination of two numpy arrays连接两个 numpy 数组的每一行组合
【发布时间】：2017-05-26 04:09:36
【问题描述】：

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

【问题讨论】：

标签： python performance numpy vectorization

【解决方案1】：

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

【讨论】：

很漂亮！
谢谢，工作得很好，并且有预先分配输出数组的附带好处，所以确保我有预先的空间。
@DaveChallis 这就是整个想法。对于大型数组，就内存/性能效率而言，任何堆叠操作都只是惩罚。
@Divakar，漂亮的解决方案！如果您有时间，能否给我一个提示，如何使用 broadcast_to 函数而不是强制显式广播连接。 github.com/numpy/numpy/issues/2115 有一个讨论，但都是基于想要加入一个 (n,1) 数组而不是整行，并且它实际上并没有使用 broadcast_to。例如，我很难将其适应于一个数组为 (3,3) 而另一个数组为 (4,2) 的情况[寻找形状为 (12,5) 的答案]。您在这里的解决方案适用于此，但希望更好地掌握 broadcast_to()。

【解决方案2】：

参考：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)

输出：

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

【讨论】：