NumPy：沿特定轴进行外部求和

Question

我正在寻找一种有效的方法来对2D矩阵的列进行外部求和。

假设A是一个(2,3)的矩阵，我想要的结果是形状为(3,2,2)的矩阵，并定义如下：

A = np.array([[a11, a12, a13],
              [a21, a22, a23]])

myWantedResult = np.array( [[[a11+a11, a11+a21],
                             [a21+a11, a21+a21]], # This is the np.add.outer(A[:,0],A[:,0])
                            [[a12+a12, a12+a22 ],
                             [a22+a12, a22+a22]], # This is the np.add.outer(A[:,1],A[:,1])
                           [[a13+a13, a13+a23],
                            [a23+a13, a23+a23]]# This is the np.add.outer(A[:,2],A[:,2])
                           ])

我已经尝试过对该列进行循环，但是由于我想在大型数组上执行此操作，因此这是非常耗时的。我正在寻找一种向量化的解决方案。

非常感谢！

- Pascal Nguyen

Answer 1

外部的加法可以使用广播来完成，将 A 加到形状为 (2, 1, 3) 的重塑版本的 A 中即可。可以通过多种方式进行重塑，包括 A[:, None, :]、A.reshape((2, 1, 3)) 或者 np.expand_dims(A, 1)。然后，A + A[:, None, :] 包含所需的结果，但形状为 (2, 2, 3)。您可以使用 np.moveaxis 来重新排列轴，使形状为 (3, 2, 2)。

In [174]: A                                                                                   
Out[174]: 
array([[  1,  10, 100],
       [-25,  50, 500]])

In [175]: B = np.moveaxis(A + A[:, None, :], -1, 0)                                           

In [176]: B.shape                                                                             
Out[176]: (3, 2, 2)

In [177]: B                                                                                   
Out[177]: 
array([[[   2,  -24],
        [ -24,  -50]],

       [[  20,   60],
        [  60,  100]],

       [[ 200,  600],
        [ 600, 1000]]])

- Warren Weckesser