相关疑难解决方法(0)

给定一个2d Numpy数组,我希望能够以最快的方式计算每一行的对角线,我现在正在使用列表理解,但我想知道它是否可以以某种方式进行矢量化？

例如,使用以下M数组:

M = np.random.rand(5, 3)


[[ 0.25891593  0.07299478  0.36586996]
 [ 0.30851087  0.37131459  0.16274825]
 [ 0.71061831  0.67718718  0.09562581]
 [ 0.71588836  0.76772047  0.15476079]
 [ 0.92985142  0.22263399  0.88027331]]

我想计算以下数组:

np.array([np.diag(row) for row in M])

array([[[ 0.25891593,  0.        ,  0.        ],
        [ 0.        ,  0.07299478,  0.        ],
        [ 0.        ,  0.        ,  0.36586996]],

       [[ 0.30851087,  0.        ,  0.        ],
        [ 0.        ,  0.37131459,  0.        ],
        [ 0.        ,  0.        ,  0.16274825]],

       [[ 0.71061831,  0.        ,  0.        ],
        [ 0.        , …