如何遍历字典时保留NaN值，将符合条件的数组元素置0？

问题背景 #

我有两个字典data_val和neg，data_val的取值为包含NaN的numpy数组。需要实现：将data_val中每个数组里小于等于neg对应key值的元素设为0，同时保留原数组中的NaN值，但当前使用的代码会把NaN转为0，需要修正。

原数据 #

data_val字典 #

{'Ancolmekar': array([0. , 0. , 0. , ..., 0. , 0. , 0.5]),
 'Cidurian': array([0., 0., 0., ..., 0., 0., 0.]),
 'Dayeuhkolot': array([nan, nan, nan, ...,  0.,  0.,  0.]),
 'Hantap': array([nan, nan, nan, ...,  0.,  0.,  0.]),
 'Kertasari': array([0., 0., 0., ..., 0., 0., 0.]),
 'Meteolembang': array([nan, nan, nan, ...,  0.,  0.,  0.]),
 'Sapan': array([0., 0., 0., ..., 0., 0., 0.])}

neg字典 #

{'Ancolmekar': -1.0,
 'Cidurian': -1.0,
 'Dayeuhkolot': -1.0,
 'Hantap': -1.0,
 'Kertasari': -1.0,
 'Meteolembang': -1.0,
 'Sapan': -1.0}

期望结果 #

{'Ancolmekar': array([0. , 0. , 0. , ..., 0. , 0. , 0.0]),
 'Cidurian': array([0., 0., 0., ..., 0., 0., 0.]),
 'Dayeuhkolot': array([nan, nan, nan, ...,  0.,  0.,  0.]),
 'Hantap': array([nan, nan, nan, ...,  0.,  0.,  0.]),
 'Kertasari': array([0., 0., 0., ..., 0., 0., 0.]),
 'Meteolembang': array([nan, nan, nan, ...,  0.,  0.,  0.]),
 'Sapan': array([0., 0., 0., ..., 0., 0., 0.])}

现有错误代码 #

for key, ch_vals in data_val.items():
    if key in neg:
        data_val[key] = [1 if val <= neg[key] else 0 for val in ch_vals]
    else:
        data_val[key] = [-1] * len(ch_vals)

错误原因 #

1. 列表推导式无法识别NaN：NaN <= 任意值的布尔结果为False，导致NaN被归入else分支转为0
2. 逻辑与需求相反：需求是将<=neg值的元素设为0，但代码里把这类元素设为1，其余设为0
3. 将numpy数组转为普通列表，破坏了原数据的NaN存储机制

修正后的代码 #

import numpy as np

for key, ch_vals in data_val.items():
    if key in neg:
        # 复制原数组，避免直接修改原始数据（可选，根据业务需求调整）
        processed_arr = ch_vals.copy()
        # 构建掩码：非NaN 且 元素小于等于neg对应值
        mask = np.logical_and(~np.isnan(processed_arr), processed_arr <= neg[key])
        # 对符合条件的元素设为0
        processed_arr[mask] = 0
        data_val[key] = processed_arr
    else:
        # 生成与原数组同类型的-1数组
        data_val[key] = np.full(len(ch_vals), -1, dtype=ch_vals.dtype)

代码说明 #

• np.isnan(processed_arr)：定位数组中的NaN位置，~取反得到非NaN元素的掩码
• np.logical_and：组合两个判断条件，确保只处理非NaN且符合数值要求的元素
• 保留numpy数组类型，天然支持NaN存储，不会丢失原数据中的空值
• 修正了原代码的逻辑错误，严格匹配需求：将<=neg对应值的元素设为0

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