如何用NumPy操作替代显式循环生成指定对角带型数组？

嘿，这个需求太适合用NumPy的向量化操作搞定了，完全能把显式循环替换掉，给你分享几种简洁高效的实现方式：

方法1：np.eye + repeat（最简洁直观） #

这种方法直接基于单位矩阵扩展，代码短到离谱：

import numpy as np

nrows, ncols = 5, 3
d = np.eye(nrows, dtype=np.uint8).repeat(ncols, axis=1)
print(d)

原理：np.eye(nrows)先生成一个nrows×nrows的单位矩阵（对角线上是1，其余全0），接着.repeat(ncols, axis=1)把每一列重复ncols次，这样每行的对角位置就被扩展成了连续的ncols个1，完美匹配目标输出。

方法2：广播索引赋值（灵活可控） #

如果之后你需要自定义块内的数值（比如不是全1），这种索引方法会更实用：

import numpy as np

nrows, ncols = 5, 3
d = np.zeros((nrows, nrows * ncols), dtype=np.uint8)

# 生成行索引（形状：(nrows, 1)）和对应的列索引块（形状：(nrows, ncols)）
row_indices = np.arange(nrows)[:, np.newaxis]
col_blocks = np.arange(ncols) + row_indices * ncols

# 批量赋值，无需循环
d[row_indices, col_blocks] = 1
print(d)

原理：利用NumPy的广播机制，row_indices和col_blocks会自动匹配维度，一次性定位到所有需要赋值的位置，直接完成批量操作，比循环灵活多了。

方法3：克罗内克积np.kron（数学视角实现） #

从线性代数的克罗内克积角度出发，也能快速生成目标结构：

import numpy as np

nrows, ncols = 5, 3
d = np.kron(np.eye(nrows, dtype=np.uint8), np.ones((1, ncols), dtype=np.uint8))
print(d)

原理：克罗内克积会把第一个矩阵（单位矩阵）的每个元素，替换成第二个矩阵（1×ncols的全1数组），这样对角线上的1就会被扩展成连续的ncols个1，非对角的0则扩展成ncols个0，正好得到我们想要的数组。

这几种方法都是NumPy的向量化操作，相比显式循环，在数据量较大时效率会提升非常明显，而且代码更简洁易读。

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