带 groupby 的 SimpleImputer

Question

假设如下数据集

	code	category	energy	sugars	proteins
0	01	B	936	NaN	7.8
1	02	NaN	NaN	15.0	NaN
2	03	A	1569.0	23	4.1
3	04	NaN	826	NaN	3
4	05	B	1345	22	5.1
5	06	A	NaN	17	NaN
6	10	C	826	NaN	3
7	11	C	1345	26	5.1
8	101	B	NaN	18	6.1
9	102	B	636	NaN	7.8
10	103	NaN	NaN	15.0	NaN
11	104	A	1569.0	23	4.1
12	105	C	813	NaN	3.5

考虑到category 列，我想用SimpleImputer 进行估算。

也就是说，考虑到产品的category，我想分配平均值。
如果产品没有类别，那么，我想考虑没有category的产品的平均值。

所以，完成 code 01 的糖。 我只会考虑所有sugars 的产品category B

	code	category	energy	sugars	proteins
0	01	B	936	NaN	7.8
4	05	B	1345	22	5.1
8	101	B	NaN	18	6.1
9	102	B	636	NaN	7.8

我做了类似的事情，如下所示。但我需要使用 SimpleImputer。
为了澄清，在下面的例子中，我用列的平均值完成了没有category的NaN。

for col in df.columns:
    if df[col].dtypes == "float64":
        df.loc[df[col].isna() & df["category"].notnull(), col] = df["categories"].map(df.groupby("category")[col].mean())
        df[col].fillna(df[col].mean(), inplace=True)

Answer 1

恐怕你不能只使用SimpleImputer来解决这类问题（至少据我所知）。

但是，您可以使用 scikit-learn 非常灵活的类 BaseEstimator 和 TransformerMixin 创建自定义的 Imputer 类。

一个非常基本的类应该是这样的：

from sklearn.base import BaseEstimator, TransformerMixin

class WithinGroupMeanImputer(BaseEstimator, TransformerMixin):
    def __init__(self, group_var):
        self.group_var = group_var
    
    def fit(self, X, y=None):
        return self
        
    def transform(self, X):
        # the copy leaves the original dataframe intact
        X_ = X.copy()
        for col in X_.columns:
            if X_[col].dtypes == 'float64':
                X_.loc[(X[col].isna()) & X_[self.group_var].notna(), col] = X_[self.group_var].map(X_.groupby(self.group_var)[col].mean())
                X_[col] = X_[col].fillna(X_[col].mean())
        return X_

---

在您的示例数据集上：

imp = WithinGroupMeanImputer(group_var='category')

imp.fit(df)

imp.transform(df)

   code category       energy     sugars  proteins
0    01        B   936.000000  20.000000  7.800000
1    02     None  1127.848485  15.000000  4.881818
2    03        A  1569.000000  23.000000  4.100000
3    04     None   826.000000  20.916667  3.000000
4    05        B  1345.000000  22.000000  5.100000
5    06        A  1569.000000  17.000000  4.100000
6    10        C   826.000000  26.000000  3.000000
7    11        C  1345.000000  26.000000  5.100000
8   101        B   972.333333  18.000000  6.100000
9   102        B   636.000000  20.000000  7.800000
10  103     None  1127.848485  15.000000  4.881818
11  104        A  1569.000000  23.000000  4.100000
12  105        C   813.000000  26.000000  3.500000

---

原始数据：

import pandas as pd

df = pd.DataFrame({
    'code': ['01', '02', '03', '04', '05', '06', '10', '11', '101', '102', '103', '104', '105'],
    'category': ['B', None, 'A', None, 'B', 'A', 'C', 'C', 'B', 'B', None, 'A', 'C'],
    'energy': [936, None, 1569, 826, 1345, None, 826, 1345, None, 636, None, 1569, 813],
    'sugars': [None, 15, 23, None, 22, 17, None, 26, 18, None, 15, 23, None],
    'proteins': [7.8, None, 4.1, 3, 5.1, None, 3, 5.1, 6.1, 7.8, None, 4.1, 3.5]
})