numpy的：计算对角线的二维数组的每一行的最快方法 [英] Numpy: Fastest way of computing diagonal for each row of a 2d array

问题描述

，我想能够计算对角线在最快的方式可能的每一行，我现在使用列表COM prehension但我想知道，如果它可以不知何故矢量化？

Given a 2d Numpy array, I would like to be able to compute the diagonal for each row in the fastest way possible, I'm right now using a list comprehension but I'm wondering if it can be vectorised somehow?

例如使用下面的 M 的数组：

For example using the following M array:

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]]

我想计算下面的数组：

I would like to compute the following array:

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]]])

推荐答案

尽管@ajcr的很好的答案，更快的替代方案，可以用花哨的索引（在numpy的1.9.0测试）实现

Despite the good answer of @ajcr, a much faster alternative can be achieved with fancy indexing (tested in NumPy 1.9.0):

import numpy as np

def sol0(M):
    return np.eye(M.shape[1]) * M[:,np.newaxis,:]

def sol1(M):
    b = np.zeros((M.shape[0], M.shape[1], M.shape[1]))
    diag = np.arange(M.shape[1])
    b[:, diag, diag] = M
    return b

在这里的时间显示，这是近4倍速度快：

where the timing shows this is approximately 4X faster:

M = np.random.random((1000, 3))
%timeit sol0(M)
#10000 loops, best of 3: 111 µs per loop
%timeit sol1(M)
#10000 loops, best of 3: 23.8 µs per loop

这篇关于numpy的：计算对角线的二维数组的每一行的最快方法的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！