python - Numpy:为二维数组的每一行计算对角线的最快方法

Question

给定一个 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.        ,  0.67718718,  0.        ],
        [ 0.        ,  0.        ,  0.09562581]],

       [[ 0.71588836,  0.        ,  0.        ],
        [ 0.        ,  0.76772047,  0.        ],
        [ 0.        ,  0.        ,  0.15476079]],

       [[ 0.92985142,  0.        ,  0.        ],
        [ 0.        ,  0.22263399,  0.        ],
        [ 0.        ,  0.        ,  0.88027331]]])

Answer 1

这是使用 np.eye(3)(3x3 恒等数组)的逐元素乘法和稍微 reshape 的 M 的一种方法:

>>> M = np.random.rand(5, 3)
>>> np.eye(3) * M[:,np.newaxis,:]
array([[[ 0.42527357,  0.        ,  0.        ],
        [ 0.        ,  0.17557419,  0.        ],
        [ 0.        ,  0.        ,  0.61920924]],

       [[ 0.04991268,  0.        ,  0.        ],
        [ 0.        ,  0.74000307,  0.        ],
        [ 0.        ,  0.        ,  0.34541354]],

       [[ 0.71464307,  0.        ,  0.        ],
        [ 0.        ,  0.11878955,  0.        ],
        [ 0.        ,  0.        ,  0.65411844]],

       [[ 0.01699954,  0.        ,  0.        ],
        [ 0.        ,  0.39927673,  0.        ],
        [ 0.        ,  0.        ,  0.14378892]],

       [[ 0.5209439 ,  0.        ,  0.        ],
        [ 0.        ,  0.34520876,  0.        ],
        [ 0.        ,  0.        ,  0.53862677]]])

(“ reshape M” 的意思是 M 的行沿 z 轴而不是横跨 y 轴面向外，给 M 形状 (5, 1, 3).)