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],[7,8]]) #shape (4,2)

ans = np.copy(x)
for i in range(x.shape[1]):
    ans[:,i] += y

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

我怎样才能沿着(例如)第二轴传播这个？更一般地说，我怎样才能沿着任意轴进行广播？

Answer 1

看两个数组的形状:

>>> x.shape
(4, 3, 2)
>>> y.shape
(4, 2)

您会看到添加需要沿此处的第 0 个和最后一个轴广播。一个简单的选择是

>>> x + y[:, None, :] 
array([[[ 2,  4],
        [ 4,  6],
        [ 6,  8]],

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

       [[ 8, 10],
        [10, 12],
        [12, 14]],

       [[16,  8],
        [ 8, 10],
        [10, 12]]])

在哪里，

>>> y[:, None, :].shape
(4, 1, 2)

这实际上只是改变了 y 的步幅，因此可以广播添加。

---

更好的是，按照 hpaulj 在评论中的建议使用 np.expand_dims，这将添加一个额外的倒数第二个维度，因此您可以这样做

>>> x + np.expand_dims(y, 1)
array([[[ 2,  4],
        [ 4,  6],
        [ 6,  8]],

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

       [[ 8, 10],
        [10, 12],
        [12, 14]],

       [[16,  8],
        [ 8, 10],
        [10, 12]]])