如何在NumPy中创建对角维度？如何高效构造特定形状的NumPy数组？

问题1：是否存在在NumPy中创建对角维度的方法？ #

当然有。NumPy针对不同维度的对角结构需求，提供了多种实现方式：

• 若要将一维数组转为二维对角数组，直接调用np.diag(arr)即可。
• 若要构造更高维度的对角结构（比如从二维扩展到三维），可以通过维度扩展、索引赋值或矩阵运算实现，这也正好对应第二个问题的场景。

问题2：从形状为(M, N)的输入数组ar构造(M, N, N)对角数组的最高效方法？ #

要求最终数组res满足：res[m][n][n] = ar[m][n]，其余位置为0。以下是几种高效实现方案：

方案1：零数组初始化+对角线索引赋值（性能最优） #

直接创建目标形状的零数组，通过广播索引一次性定位所有对角线位置并赋值，内存访问效率极高，计算开销最小：

import numpy as np

M, N = ar.shape
res = np.zeros((M, N, N), dtype=ar.dtype)
# 利用广播索引定位所有(m, n, n)位置
res[np.arange(M)[:, None], np.arange(N), np.arange(N)] = ar

方案2：使用np.einsum（简洁高效） #

爱因斯坦求和符号可以用非常简洁的语法构造目标结构，底层经过NumPy优化，性能接近方案1：

res = np.einsum('mn,np->mnp', ar, np.eye(N, dtype=ar.dtype))

原理是将ar的每一行与N阶单位矩阵做对应运算，生成每个M维度对应的N×N对角矩阵。

方案3：逐行生成对角矩阵（直观但性能稍弱） #

对ar的每一行单独生成对角矩阵，再堆叠成三维数组，适合小数据量场景：

res = np.array([np.diag(row) for row in ar])

这种方式本质是Python层面的循环，数据量较大时性能不如前两种方案。

总结：追求极致性能选方案1；看重代码简洁性选方案2。

内容的提问来源于stack exchange，提问作者extremeaxe5