如何让MICE仅为布尔型变量生成0和1的缺失值补全结果？

问题原因 #

IterativeImputer 默认使用贝叶斯岭回归这类回归模型进行迭代补全，输出是连续数值。即便设置了 min_value 和 max_value 限制范围，结果仍会是0-1之间的小数，不会自动离散为布尔型的0/1。

解决方案 #

方案1：补全后对布尔列强制离散化

先标记出数据集中的布尔型列（包括用0/1数值存储的布尔列），补全后直接将这些列的数值按阈值截断为0/1，再转回布尔类型。

修改后的代码示例：

def do_MICE_imputation(train, test):
    import numpy as np
    import pandas as pd
    from sklearn.experimental import enable_iterative_imputer
    from sklearn.impute import IterativeImputer

    # 标记布尔型列：包括原生bool类型，以及值仅为0/1的数值列
    boolean_cols = []
    for col in train.columns:
        unique_vals = train[col].dropna().unique()
        if train[col].dtype == 'bool' or set(unique_vals).issubset({0, 1}):
            boolean_cols.append(col)
    
    # 记录各列的极值
    min_value , max_value = {}, {}
    for col in train.columns:
        min_value[col] = np.min(train[col])
        max_value[col] = np.max(train[col])
        
    # 初始化MICE补全器
    imp = IterativeImputer(
        random_state=0, 
        min_value=list(min_value.values()),
        max_value=list(max_value.values()), 
        max_iter=5, 
        verbose=10,
        initial_strategy='most_frequent'
    )
    imp.fit(train)
    
    # 补全数据
    train_imputed = pd.DataFrame(imp.transform(train), columns=train.columns)
    test_imputed = pd.DataFrame(imp.transform(test), columns=test.columns)
    
    # 对布尔列进行离散化处理
    for col in boolean_cols:
        # 按0.5阈值截断为0/1，再转回布尔类型（或保持int类型）
        train_imputed[col] = np.where(train_imputed[col] >= 0.5, 1, 0).astype(bool)
        test_imputed[col] = np.where(test_imputed[col] >= 0.5, 1, 0).astype(bool)
    
    return train_imputed, test_imputed

方案2：针对布尔列使用分类模型补全

自定义迭代补全的估计器，让模型在处理布尔型列时使用逻辑回归（分类模型），直接输出0/1；处理数值列时仍用回归模型。这种方法更贴合布尔变量的分布，补全结果更准确。

代码示例：

def do_MICE_imputation(train, test):
    import numpy as np
    import pandas as pd
    from sklearn.experimental import enable_iterative_imputer
    from sklearn.impute import IterativeImputer
    from sklearn.linear_model import BayesianRidge, LogisticRegression
    from sklearn.base import BaseEstimator, RegressorMixin

    # 标记布尔型列
    boolean_cols = []
    for col in train.columns:
        unique_vals = train[col].dropna().unique()
        if train[col].dtype == 'bool' or set(unique_vals).issubset({0, 1}):
            boolean_cols.append(col)
    
    # 自定义估计器：根据目标列类型选择回归/分类模型
    class CustomImputeEstimator(BaseEstimator, RegressorMixin):
        def __init__(self, boolean_cols):
            self.boolean_cols = boolean_cols
            self.reg_model = BayesianRidge()
            self.clf_model = LogisticRegression(max_iter=1000)
        
        def fit(self, X, y):
            # 判断当前补全的列是否为布尔列
            if y.name in self.boolean_cols:
                self.clf_model.fit(X, y.astype(int))
                self.model = self.clf_model
            else:
                self.reg_model.fit(X, y)
                self.model = self.reg_model
            return self
        
        def predict(self, X):
            return self.model.predict(X)
    
    # 初始化MICE补全器，使用自定义估计器
    imp = IterativeImputer(
        random_state=0, 
        estimator=CustomImputeEstimator(boolean_cols=boolean_cols),
        max_iter=5, 
        verbose=10,
        initial_strategy='most_frequent'
    )
    imp.fit(train)
    
    # 补全数据（布尔列直接输出0/1）
    train_imputed = pd.DataFrame(imp.transform(train), columns=train.columns)
    test_imputed = pd.DataFrame(imp.transform(test), columns=test.columns)
    
    # 可选：将布尔列转回bool类型
    for col in boolean_cols:
        train_imputed[col] = train_imputed[col].astype(bool)
        test_imputed[col] = test_imputed[col].astype(bool)
    
    return train_imputed, test_imputed

方案对比 #

• 方案1实现简单，适合快速处理，但四舍五入/截断可能引入少量偏差；
• 方案2更符合布尔变量的统计特性，补全结果更可靠，但需要自定义模型，代码稍复杂。

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