如何用NumPy内置函数替代列表推导提取非NaN值数组？

核心思路 #

用NumPy的布尔索引替代列表推导，完全依托向量化操作提升效率，同时严格保留原顺序。

实现代码 #

import numpy as np

# 假设a是形状为(14, 3200, 2500)的输入数组
mask = ~np.isnan(a).any(axis=0)  # 生成(3200,2500)的布尔掩码，标记无NaN的位置

# 获取索引数组，形状(-1,2)，与原代码输出完全一致
idx = np.argwhere(mask)

# 提取非NaN对应的14维向量，形状(-1,14)
vals = a[:, mask].T

代码说明 #

1. 掩码计算：np.isnan(a).any(axis=0)在14维（第一个轴）上判断每个(i,j)位置是否存在NaN，取反后得到标记有效位置的布尔掩码。
2. 索引获取：np.argwhere(mask)直接返回所有有效位置的坐标数组，形状为(-1,2)，和原代码中np.stack((indices[0],indices[1])).transpose()的结果完全一致，顺序保持不变。
3. 值提取：a[:, mask]通过布尔索引直接提取所有有效位置的14维数据，得到形状为(14, N)的数组（N为有效位置总数），转置后即得到要求的(-1,14)形状。

效率优势 #

原列表推导是Python级别的循环，每个迭代都要执行索引查找；而布尔索引是NumPy底层C实现的向量化操作，速度能提升几十到上百倍，尤其适配大尺寸数组场景。

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