有没有一种方法可以强制SimpleImputer返回熊猫数据框?

时间:2020-06-04 09:47:11

标签: python pandas numpy scikit-learn

以下代码

from sklearn.impute import SimpleImputer
import pandas as pd

df = pd.DataFrame(dict(
    x=[1, 2, np.nan],
    y=[2, np.nan, 0]
))

SimpleImputer().fit_transform(df)

返回

array([[1. , 2. ],
       [2. , 1. ],
       [1.5, 0. ]])

是否有一种方法可以使用返回大熊猫数据框而不是numpy数组的imputer?为此有一个scikit-learn实现吗?

3 个答案:

答案 0 :(得分:2)

SimpleImputer返回一个转换后的数组,如docs中所述。您可以构造一个新的数据框,也可以使用来自冲床的结果就地修改该数据框,如下所示:

df[:] = SimpleImputer().fit_transform(df)

答案 1 :(得分:1)

通常,sklearn中的所有转换函数都返回numpy数组。您需要养成使用

的习惯
pd.DataFrame(SimpleImputer().fit_transform(df), columns = df.columns)

答案 2 :(得分:1)

如果您想保留列(例如,为了在后面的步骤中使用 ColumnTransformers),您可以围绕 SimpleImputer 创建一个包装器:

df = pd.DataFrame({"A": [1, 2, np.NaN], "B": [3, np.NaN, 4], "C": [np.NaN, 5, 6]})

class PandasSimpleImputer(SimpleImputer):
    """A wrapper around `SimpleImputer` to return data frames with columns.
    """

    def fit(self, X, y=None):
        self.columns = X.columns
        return super().fit(X, y)

    def transform(self, X):
        return pd.DataFrame(super().transform(X), columns=self.columns)


PandasSimpleImputer().fit_transform(df)

>>>
    A   B   C
0   1.0 3.0 5.5
1   2.0 3.5 5.0
2   1.5 4.0 6.0
相关问题