基于NumPy实现按1D计数数组填充2D数组列的高效方法

当然可以用NumPy的向量化操作替代循环，既简洁又能在处理大规模数组时大幅提升效率！核心思路是利用广播机制和布尔掩码来实现批量赋值，完全不用逐列遍历：

import numpy as np

cnt = np.array([1, 3, 2, 4])
a = np.zeros((5, 4))

# 生成行索引并转成可广播的形状
row_idx = np.arange(a.shape[0])[:, np.newaxis]
# 用布尔掩码批量赋值1
a[row_idx < cnt] = 1

print(a)

运行后会得到和你循环代码一模一样的输出：

[[1. 1. 1. 1.]
 [0. 1. 1. 1.]
 [0. 1. 0. 1.]
 [0. 0. 0. 1.]
 [0. 0. 0. 0.]]

原理拆解 #

1. 构造可广播的行索引：np.arange(a.shape[0])生成[0,1,2,3,4]的行索引，通过[:, np.newaxis]将其转换为(5,1)的二维数组，这样就能和形状为(4,)的cnt数组触发广播，最终得到一个(5,4)的矩阵。
2. 生成布尔掩码：row_idx < cnt会逐元素比较每个行索引是否小于对应列的cnt值，得到的布尔矩阵中，True的位置就是需要填充1的位置。
3. 批量赋值：直接用这个布尔掩码给数组a赋值，一步完成所有列的填充操作，全程没有Python循环的开销。

关于你尝试的np.full失败的原因 #

你用np.full((5, 4), 1, cnt)无法运行，是因为np.full的第三个参数是指定数据类型的dtype，不是用来控制填充数量的参数，所以这个用法不符合函数的定义哦。

如果你的数组规模较大，这个向量化方案的性能优势会非常明显——毕竟NumPy的底层操作是用C实现的，比Python级别的循环快得多。

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