Python 从二维矩阵的行创建对角矩阵，而不使用for循环

假设您有一个

mxn

矩阵a，并且想要创建m个对角矩阵，每个对角矩阵都来自a的行，因此具有

nxn

的形状。结果矩阵的形状应为

mxnxn

我知道一个典型的解决方案是：

result = numpy.stack([numpy.diag(A[i,:]) for i in range(A.shape[0])], axis=0)

我想知道我们是否可能不使用循环就能得到相同的结果

---

任何想法都将不胜感激

这不是最有效的方法，但类似的方法可能会奏效

import numpy as np

A = np.random.rand(10, 5)

S = np.einsum("ai,ij->aij", A, np.ones((5, 5)))
M = np.eye(5).astype(np.bool)
M = np.repeat(M[None, ...], 10, axis=0)
S[~M]=0
print(S.shape)

这似乎有效：

result = np.zeros((A.shape[0], A.shape[1], A.shape[1]), dtype=A.dtype)
result[:, range(A.shape[1]), range(A.shape[1])] = A

测试输入：

A = np.arange(24).reshape(4,6)

[[[ 0  0  0  0  0  0]
  [ 0  1  0  0  0  0]
  [ 0  0  2  0  0  0]
  [ 0  0  0  3  0  0]
  [ 0  0  0  0  4  0]
  [ 0  0  0  0  0  5]]

 [[ 6  0  0  0  0  0]
  [ 0  7  0  0  0  0]
  [ 0  0  8  0  0  0]
  [ 0  0  0  9  0  0]
  [ 0  0  0  0 10  0]
  [ 0  0  0  0  0 11]]

 [[12  0  0  0  0  0]
  [ 0 13  0  0  0  0]
  [ 0  0 14  0  0  0]
  [ 0  0  0 15  0  0]
  [ 0  0  0  0 16  0]
  [ 0  0  0  0  0 17]]

 [[18  0  0  0  0  0]
  [ 0 19  0  0  0  0]
  [ 0  0 20  0  0  0]
  [ 0  0  0 21  0  0]
  [ 0  0  0  0 22  0]
  [ 0  0  0  0  0 23]]]

打印输出（结果）：

A = np.arange(24).reshape(4,6)

[[[ 0  0  0  0  0  0]
  [ 0  1  0  0  0  0]
  [ 0  0  2  0  0  0]
  [ 0  0  0  3  0  0]
  [ 0  0  0  0  4  0]
  [ 0  0  0  0  0  5]]

 [[ 6  0  0  0  0  0]
  [ 0  7  0  0  0  0]
  [ 0  0  8  0  0  0]
  [ 0  0  0  9  0  0]
  [ 0  0  0  0 10  0]
  [ 0  0  0  0  0 11]]

 [[12  0  0  0  0  0]
  [ 0 13  0  0  0  0]
  [ 0  0 14  0  0  0]
  [ 0  0  0 15  0  0]
  [ 0  0  0  0 16  0]
  [ 0  0  0  0  0 17]]

 [[18  0  0  0  0  0]
  [ 0 19  0  0  0  0]
  [ 0  0 20  0  0  0]
  [ 0  0  0 21  0  0]
  [ 0  0  0  0 22  0]
  [ 0  0  0  0  0 23]]]

---

您可以在我们使用和的地方尝试以下方法，以获得所需的结果

# `arr` is input of size (m, n); for instance,
arr = np.arange(1, 13).reshape(3, -1)   # (3, 4)

# get column size
In [160]: _ , n = arr.shape

# fill the rows from the input, after promoting it to 3D
In [161]: res = np.eye(n, dtype=arr.dtype) * arr[:,np.newaxis]

In [162]: res.shape
Out[162]: (3, 4, 4)

---

注意：由于我们将

（n，n）

作为最后两个维度，因此也可以使用

In [171]: res = np.identity(n, dtype=arr.dtype) * arr[:,np.newaxis]

In [172]: res.shape
Out[172]: (3, 4, 4)

---

另一种方法是使用where，首先用零初始化结果数组，切掉对角元素的索引，并用输入数组填充它们