如何创建包含N个M维对角矩阵的np.array？

要生成形状为(N, M, M)的numpy数组，其中每个子数组是M×M的对角矩阵（主对角线为随机元素，其余为0），可以用这两种方法：

方法1：向量化赋值（高效推荐） #

这种方法用numpy的向量化操作，避免循环，在N或M较大时性能更好：

import numpy as np

N = 2
M = 3

# 生成N组、每组M个0-1之间的随机数，对应每个对角矩阵的主对角线元素
diag_vals = np.random.rand(N, M)
# 创建全零的目标数组，形状是(N, M, M)
result = np.zeros((N, M, M))
# 通过索引把随机值放到每个矩阵的主对角线上
result[:, np.arange(M), np.arange(M)] = diag_vals

运行后result的形状就是(2,3,3)，每个子矩阵都是符合要求的对角矩阵。

方法2：列表推导式（直观易读） #

如果N不大，用列表推导式逐个生成对角矩阵再组合成数组，代码更直观：

import numpy as np

N = 2
M = 3

# 循环N次，每次生成一个M维随机对角矩阵，再转成numpy数组
result = np.array([np.diag(np.random.rand(M)) for _ in range(N)])

两种方法都能得到你想要的结构，主对角线元素都是随机生成的。

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