如何用Numpy高效填充5×3×3数组中每个3×3块的对角线？

嘿，这个需求用Numpy完全可以高效实现，不用写繁琐的嵌套循环，我给你分享两种实用的方法：

方法一：直接利用高级索引赋值（适合已有全零数组的场景） #

如果已经初始化了全零的5×3×3数组，我们可以通过Numpy的广播式高级索引直接定位所有子块的对角线位置，一次性完成赋值：

import numpy as np

# 初始化目标形状的全零数组
arr = np.zeros((5, 3, 3))

# 准备每个子块的对角线数值（替换成你需要的数值即可）
diag_values = np.array([
    [0.07735655, 0.11476396, 0.09903619],
    [0.1923885, 0.03063454, 0.06028193],
    [0.06566275, 0.03151423, 0.04042383],
    [0.07950743, 0.03250461, 0.0448308],
    [0.10879917, 0.04700161, 0.03924387]
])

# 定位所有对角线位置并赋值
arr[np.arange(5)[:, None], np.arange(3), np.arange(3)] = diag_values

原理说明：

• np.arange(5)[:, None] 把一维的[0,1,2,3,4]转变成形状为(5,1)的数组，这样能和np.arange(3)（形状(3,)）进行广播，生成所有子块的对角线索引：(0,0,0)、(0,1,1)、(0,2,2)、(1,0,0)...直到(4,2,2)。
• 这种方式是纯矢量操作，没有Python层面的循环，效率拉满。

方法二：直接构造对角块再堆叠（适合从对角值生成目标数组） #

如果还没初始化数组，直接从对角线数值生成目标数组会更简洁：

import numpy as np

# 同样准备对角线数值
diag_values = np.array([
    [0.07735655, 0.11476396, 0.09903619],
    [0.1923885, 0.03063454, 0.06028193],
    [0.06566275, 0.03151423, 0.04042383],
    [0.07950743, 0.03250461, 0.0448308],
    [0.10879917, 0.04700161, 0.03924387]
])

# 对每一行对角值生成3×3对角矩阵，再堆叠成5×3×3数组
arr = np.stack([np.diag(row) for row in diag_values])

原理说明：

• np.diag(row) 会把输入的一维数组转换成对应的对角矩阵，比如输入[a,b,c]就得到[[a,0,0],[0,b,0],[0,0,c]]。
• np.stack() 把5个3×3的对角矩阵沿第一个维度堆叠，直接得到目标形状的数组，代码可读性非常好。

扩展：随机生成对角线数值 #

如果你的对角线数值是随机生成的（比如用于测试），可以用以下方式快速实现：

# 生成5组、每组3个随机数作为对角线值
random_diag = np.random.rand(5, 3)
# 方法一赋值
arr = np.zeros((5,3,3))
arr[np.arange(5)[:, None], np.arange(3), np.arange(3)] = random_diag

# 或者方法二直接生成
arr = np.stack([np.diag(np.random.rand(3)) for _ in range(5)])

内容的提问来源于stack exchange，提问作者Ahmed Junaid Khalid