代码之家  ›  专栏  ›  技术社区  ›  Him

沿任意轴的numpy广播加法

  •  1
  • Him  · 技术社区  · 7 年前

    非矢量化解决方案:

    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)
    

    如何沿(例如)第二个轴广播?一般来说,我怎样才能使这个广播沿任意轴?

    1 回复  |  直到 7 年前
        1
  •  7
  •   cs95 abhishek58g    7 年前

    请看两个阵列的形状:

    >>> 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 所以这个加法可以广播。


    最好还是用 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]]])