Python 如何在numpy中使用2d数组构造对角数组？

使用

np.diag

我能够构造一个二维数组，其中输入一个一维数组，然后在对角线上返回。但是如果我有n-D数组作为输入，如何做同样的事情呢

这很有效

foo = np.random.randint(2, size=(36))
print foo
print np.diag(foo)
[1 1 1 1 1 1 0 1 0 0 1 1 0 1 1 1 1 1 0 0 0 0 0 0 0 1 0 1 0 1 0 0 0 1 1 0]
[[1 0 0 ..., 0 0 0]
 [0 1 0 ..., 0 0 0]
 [0 0 1 ..., 0 0 0]
 ..., 
 [0 0 0 ..., 1 0 0]
 [0 0 0 ..., 0 1 0]
 [0 0 0 ..., 0 0 0]]

这不会

foo = np.random.randint(2, size=(2,36))
print foo
[[1 1 0 1 1 0 1 1 1 1 1 0 0 0 0 0 1 0 0 1 0 0 0 0 1 0 0 0 1 0 0 0 0 1 0 0]
 [0 1 1 1 0 0 1 1 1 1 0 1 0 1 0 1 0 0 0 0 1 0 1 1 1 1 0 0 1 0 1 1 1 1 0 1]]
do_something(foo)

应该回来

array([[[ 1.,  0.,  0., ...,  0.,  0.,  0.],
        [ 0.,  1.,  0., ...,  0.,  0.,  0.],
        [ 0.,  0.,  0., ...,  0.,  0.,  0.],
        ..., 
        [ 0.,  0.,  0., ...,  1.,  0.,  0.],
        [ 0.,  0.,  0., ...,  0.,  0.,  0.],
        [ 0.,  0.,  0., ...,  0.,  0.,  0.]],

       [[ 0.,  0.,  0., ...,  0.,  0.,  0.],
        [ 0.,  1.,  0., ...,  0.,  0.,  0.],
        [ 0.,  0.,  1., ...,  0.,  0.,  0.],
        ..., 
        [ 0.,  0.,  0., ...,  1.,  0.,  0.],
        [ 0.,  0.,  0., ...,  0.,  0.,  0.],
        [ 0.,  0.,  0., ...,  0.,  0.,  1.]]])

编辑

一些测试基于艾伦和ajcr在这篇文章中的回答以及Saulo Castro和jaime对ajcr的回答。和往常一样，这完全取决于你的投入。我的输入通常具有以下形状：

M = np.random.randint(2, size=(1000, 36))

其功能如下：

def Alan(M):
    M = np.asarray(M)
    depth, size = M.shape
    x = np.zeros((depth,size,size))
    for i in range(depth):
        x[i].flat[slice(0,None,1+size)] = M[i]
    return x

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

def Saulo(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

def jaime(M):
    b = np.zeros((M.shape[0], M.shape[1]*M.shape[1]))
    b[:, ::M.shape[1]+1] = M
    return b.reshape(M.shape[0], M.shape[1], M.shape[1])

我的结果如下

%timeit Alan(M)
100 loops, best of 3: 2.22 ms per loop    
%timeit ajcr(M)
100 loops, best of 3: 5.1 ms per loop    
%timeit Saulo(M)
100 loops, best of 3: 4.33 ms per loop    
%timeit jaime(M)
100 loops, best of 3: 2.07 ms per loop

---

在纯NumPy中，一种简单的方法是执行以下数组乘法：

np.eye(foo.shape[1]) * foo[:, np.newaxis]

其中

foo

是对角线的二维数组

这将NxN标识数组与

foo

的每一行相乘，以生成所需的3D矩阵

由于此方法中的语法相当简单，因此可以轻松地将其扩展到更高的维度。例如：

>>> foo = np.array([[0, 1], [1, 1]])
>>> d = np.eye(foo.shape[1]) * foo[:, np.newaxis] # 2D to 3D
>>> d
array([[[ 0.,  0.],
        [ 0.,  1.]],

       [[ 1.,  0.],
        [ 0.,  1.]]])

>>> np.eye(d.shape[1]) * d[:, :, np.newaxis] # 3D to 4D
array([[[[ 0.,  0.],
         [ 0.,  0.]],

        [[ 0.,  0.],
         [ 0.,  1.]]],

       [[[ 1.,  0.],
         [ 0.,  0.]],

        [[ 0.,  0.],
         [ 0.,  1.]]]])

---

可能是相关的；它还显示了一种从2D数组中导出所需对角矩阵的更快（但更详细）的方法。

我试图做一些奇特的索引，并使用

numpy。diag_Indexs_来自
，没想到我会在这里使用for循环。这种方法将创建额外的标识矩阵，并将执行许多不必要的乘法。如果
foo
很大，这可能会很昂贵。当然，这种方法的主要优点是简单的语法允许它相对容易地扩展到更高的维度。对于较小的数组，此方法仍然非常快速，但对于较大的数组，使用索引的方法肯定会更有效（特别是如果使用NumPy 1.9）。
>>> foo = np.array([[0, 1], [1, 1]])
>>> d = np.eye(foo.shape[1]) * foo[:, np.newaxis] # 2D to 3D
>>> d
array([[[ 0.,  0.],
        [ 0.,  1.]],

       [[ 1.,  0.],
        [ 0.,  1.]]])

>>> np.eye(d.shape[1]) * d[:, :, np.newaxis] # 3D to 4D
array([[[[ 0.,  0.],
         [ 0.,  0.]],

        [[ 0.,  0.],
         [ 0.,  1.]]],

       [[[ 1.,  0.],
         [ 0.,  0.]],

        [[ 0.,  0.],
         [ 0.,  1.]]]])