如何以对角方式填充ndarray？求shape为(a,b,c,b)数组的更优实现

你的代码逻辑没问题，这里提供两种更简洁高效的实现思路：

方法一：利用广播直接生成（推荐） #

借助np.eye生成b×b的单位矩阵，再通过扩展维度配合广播机制，直接得到目标形状的数组，无需先创建全零数组再赋值：

import numpy as np

a, b, c = 2, 3, 5
# 将单位矩阵扩展为(1, b, 1, b)的形状，再广播到(a, b, c, b)
arr = np.broadcast_to(np.eye(b)[None, :, None, :], (a, b, c, b))

这种方式的优势是高效省内存：broadcast_to不会额外复制数据（仅创建视图），如果需要可修改的数组，后续加.copy()即可。

方法二：简化原代码的索引写法 #

保留你原有的逻辑，去掉多余的变量定义，直接用np.arange(b)作为索引，代码更紧凑：

import numpy as np

a, b, c = 2, 3, 5
arr = np.zeros((a, b, c, b))
arr[:, np.arange(b), :, np.arange(b)] = 1

两种方式都能实现需求，第一种在性能和代码简洁性上更优，适合大规模数据场景。

内容的提问来源于stack exchange，提问作者吴慈霆