Numpy添加所有向量组合

Question

什么是最有效的方法来添加两个向量数组的所有组合？例如，我想要的是：

a = np.array([[1,2,3], [4,5,6]])
b = np.array([[7,8,9], [10,11,12]])
[ai + bj for ai in a for bj in b]

给出

[array([ 8, 10, 12]),
 array([11, 13, 15]),
 array([11, 13, 15]),
 array([14, 16, 18])]

它是一个带有向量的网格，而不是主数据类型。

我尝试过显式地构造meshgrid结果，这比列表理解更快：

a_tile = np.tile(a, (2, 1))

array([[1, 2, 3],
       [4, 5, 6],
       [1, 2, 3],
       [4, 5, 6]])

b_repeat = np.repeat(b, 2, axis=0)

array([[ 7,  8,  9],
       [ 7,  8,  9],
       [10, 11, 12],
       [10, 11, 12]])

a_tile + b_repeat

array([[ 8, 10, 12],
       [11, 13, 15],
       [11, 13, 15],
       [14, 16, 18]])

这有效率吗？我在寻找一种广播数组的方法，这样就不会显式地构造网格了。

Answer 1

您可以尝试以下方法：

import numpy as np

a = np.array([[1,2,3], [4,5,6]])
b = np.array([[7,8,9], [10,11,12]])

(a[..., None] + b.T).transpose(0, 2, 1).reshape(-1, 3)

它规定：

array([[ 8, 10, 12],
       [11, 13, 15],
       [11, 13, 15],
       [14, 16, 18]])

Answer 2

您可以使用至广播数组。

N = 2 #number of repeats
your_req_array = np.broadcast_to(b.T, (N,b.shape[1], b.shape[0])).transpose(2,0,1).reshape(-1,b.shape[1]) + np.broadcast_to(a, (N,a.shape[0], a.shape[1])).reshape(-1,b.shape[1])