python - numpy 沿第一个轴添加

Question

我想通过简单地沿第一个轴执行相同的加法来添加两个具有不同维度的数组。

非矢量化解决方案：

x = np.array([[[1,2],[3,4],[5,6]],[[7,8],[9,0],[1,2]],[[3,4],[5,6],[7,8]],[[9,0],[1,2],[3,4]]]) #shape (4,3,2)
y = np.array([[1,2],[3,4],[5,6]]) #shape (3,2)

ans = np.empty(x.shape)
for i in range(x.shape[0]):
    ans[i] = x[i] + y

print(ans) #shape (4,3,2)

我怎样才能适当地进行这个广播？

标签： pythonnumpy

Answer 1

由于广播 [numpy-doc]，您可以简单地使用：

x + y

所以在这里我们计算索引i,j,k处的元素：

x _ijk +y _jk

这给出了：

>>> x + y
array([[[ 2,  4],
        [ 6,  8],
        [10, 12]],

       [[ 8, 10],
        [12,  4],
        [ 6,  8]],

       [[ 4,  6],
        [ 8, 10],
        [12, 14]],

       [[10,  2],
        [ 4,  6],
        [ 8, 10]]])
>>> (x + y).shape
(4, 3, 2)

如果将两个数组加在一起，例如第一个数组具有三个维度，第二个数组具有三个维度，并且第一个左侧数组的最后两个维度等于右侧数组的维度，则右侧的数组为“广播”。这意味着它被视为一个三维数组，其中每个子数组等于右侧的数组。

您还可以在任意位置“引入”额外维度，y例如在此答案中“广播”特定维度。