如何基于类别列空值，为指定One-Hot编码列批量填充NaN？

问题说明 #

需要实现两个逻辑：

• 当category_with_order列取值为NaN时，将cat1 = ['a','b','c','d']对应的所有独热编码列整行设为NaN
• 当category_without_order列取值为NaN时，将cat2 = ['drinks', 'meat', 'plant', 'salad']对应的所有独热编码列整行设为NaN

修正并完善后的代码 #

import pandas as pd
import numpy as np
from sklearn.preprocessing import OneHotEncoder
from sklearn.impute import KNNImputer

# 初始化数据
data = {
    'category_with_order': ['a', 'b', 'c','d',np.nan],
    'category_without_order': ['plant',np.nan,'salad','meat', 'drinks'],
    'sugar': ['1',np.nan, '2', '2',np.nan],
    'salt': ['1',np.nan, '2', '1',np.nan]
}

# 修正变量名笔误：data1改为data
df = pd.DataFrame(data)

# 对分类列做独热编码
ohe = OneHotEncoder()
feature_array = ohe.fit_transform(df[["category_with_order","category_without_order"]]).toarray()
features_labels = ohe.categories_
feature_labels = np.hstack([i.ravel() for i in features_labels])
features = pd.DataFrame(feature_array, columns = feature_labels)

# 定义需要处理的类别组
cat1 = ['a','b','c','d']
cat2 = ['drinks', 'meat', 'plant', 'salad']

# 处理第一个条件：category_with_order为NaN时，cat1对应列置空
mask1 = df['category_with_order'].isna()
features.loc[mask1, cat1] = np.nan

# 处理第二个条件：category_without_order为NaN时，cat2对应列置空
mask2 = df['category_without_order'].isna()
features.loc[mask2, cat2] = np.nan

# 合并原数据和处理后的独热编码列，修正df1为df并补全concat语法
df_new = pd.concat([df.reset_index(drop=True), features.reset_index(drop=True)], axis=1)

# 后续KNN缺失值填充（可选）
knn = KNNImputer(n_neighbors=1, add_indicator=True)
df_imputed = pd.DataFrame(knn.fit_transform(df_new), columns=df_new.columns)

print(df_imputed)

关键逻辑解释 #

1. 修正语法错误：初始代码里的data1、df1是笔误，统一修正为data和df；补全pd.concat的闭合括号并指定axis=1实现按列合并
2. 生成条件掩码：用isna()生成布尔掩码，精准定位需要置空的行
3. 批量置空操作：通过df.loc[掩码, 列名列表]直接对目标列的指定行批量赋值为NaN，高效完成需求
4. 后续填充支持：如果需要继续做缺失值填充，直接用KNNImputer处理合并后的完整数据集即可

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