如何用Numpy临时移除数组NaN值并实现后续还原？

问题原因 #

你之前用np.insert失败的核心原因是：原数组中NaN的索引基于原数组长度，但移除NaN得到no_nan后数组长度变短。插入第一个NaN会让后续索引偏移，导致NaN被插入错误位置。

正确矢量化实现 #

步骤1：提取非NaN数组并记录位置掩码

用布尔掩码标记原数组中非NaN的位置，这是最直接的矢量化方式：

import numpy as np

# 原数组
array = np.array([np.nan, 1., 2., 3., 4., 5., np.nan, 6., 7., 8., 9.])
# 生成掩码：True = 非NaN，False = NaN
mask = ~np.isnan(array)
# 提取不含NaN的数组
no_nan = array[mask]

步骤2：还原原数组

创建和原数组形状一致的全NaN数组，通过掩码将no_nan的值填充回对应位置：

# 初始化全NaN数组，形状与原数组一致
original_array = np.full_like(array, np.nan)
# 将非NaN值填充回原位置
original_array[mask] = no_nan

执行后得到的original_array与原数组完全一致：

array([nan,  1.,  2.,  3.,  4.,  5., nan,  6.,  7.,  8.,  9.])

另一种实现方式（索引映射） #

如果需要记录具体索引而非掩码，也可以通过有效索引映射还原：

# 记录原数组中非NaN的索引
valid_indices = np.where(mask)[0]
# 初始化全NaN数组，将no_nan赋值给对应有效索引
original_array = np.full_like(array, np.nan)
original_array[valid_indices] = no_nan

效果和掩码方法完全一致，只是逻辑写法不同。

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