Python使用MICE填充缺失值的简单实现

#coding:utf-8
import os
import pandas as pd
import numpy as np
# from impyute.imputation.cs import mice
from fancyimpute import MICE

"""
使用MICE对数据集缺失值进行填充

1.分别读取train/test数据
2.合并训练测试集
3.使用MICE对缺失值进行填充
4.填充后数据重新划分为训练集/测试集
"""
#1.分别读取训练/测试数据为dataframe格式
train = pd.read_csv(r"D:\data\train.csv",encoding="utf-8")
test = pd.read_csv(r"D:\data\test.csv",encoding="utf-8")

#2.训练/测试数据集合并
df = train.append(test)

#3.使用MICE填充缺失值
mice = MICE()
new_df = mice.complete(np.asarray(df.values, dtype=float))

#4.将填充后的数据集重新划分为训练/测试集
columns = train.columns.tolist()
train_shape = train.shape[0]
df_train = pd.DataFrame(new_df[:train_shape,:],columns=columns)
df_test = pd.DataFrame(new_df[train_shape:,:],columns=columns)