基于NumPy向量化计算：从数组a构造二维矩阵b的高效方法

你这个需求用循环虽然能跑，但确实不够高效，尤其是当数组a的长度或者最大值比较大的时候，咱们直接用NumPy的向量化操作来搞定，完全不用写循环，速度快很多。

先明确需求：矩阵b的行数是max(a)+1，每一列i的元素是从0到a[i]（包含a[i]），行号超过a[i]的位置补NaN。

直接上高效的向量化实现代码：

import numpy as np

a = np.array([2, 3])  # 先把列表转成NumPy数组，方便广播
rows = np.max(a) + 1
# 生成行索引的列向量，用来和每一列的范围值比较
row_idx = np.arange(rows)[:, np.newaxis]
# 用where根据条件填充值，自动广播实现逐列判断
b = np.where(row_idx <= a, row_idx, np.nan)

运行后得到的结果正好符合你的期望：

array([[ 0.,  0.],
       [ 1.,  1.],
       [ 2.,  2.],
       [nan,  3.]])

为什么这个方法高效？ #

• 完全利用NumPy的广播机制：row_idx是(rows,1)的列向量，a是(1,cols)的一维数组，两者比较时会自动扩展成(rows,cols)的布尔矩阵，不用手动循环每一列。
• np.where是向量化的条件选择函数，底层是C实现的，比Python层面的循环快几个数量级，尤其是当数据规模变大时，优势特别明显。

对比你原来的循环代码，这个方法省去了手动初始化矩阵和循环赋值的步骤，代码更简洁，执行效率也更高。

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