For the dataset I’m working with, I want to impute null values with medians grouped by country. I have created a grouped table (median_data in the code below), which contains all the median values by country.
I need to do this median calculation and imputation in separate steps since the end goal is to create an object with ‘fit’ and ‘transform’ methods, so that I can calculate medians based only on the train data, and impute to the test data.
Here’s dummy data I’m working with:
data = [['A', 10, 20, np.nan, np.nan, 50, 30], ['A', 2, 1, 5, np.nan, 34, 35], ['A', 13, 212, 3, 6, np.nan, 37],
['B', 120, 230, 53, np.nan, 63, 23], ['B', 22, 115, 15, 61, 4, 15], ['B', np.nan, 22, 12, np.nan, np.nan, 31],
['C', 105, 120, np.nan, 22, 520, 3], ['C', 26, 11, 15, np.nan, 34, 3], ['C', 13, np.nan, 13, 234, np.nan, 10],
['D', 101, 220, 654, 143, 634, 123], ['D', 32, 21, 61, 24, np.nan, 32], ['D', 11, 72, 23, np.nan, 534, 30]
]
df = pd.DataFrame(data, columns=['Country','col1','col2','col3','col4','col5','col6'])
median_data = df.groupby('Country').median().reset_index()
Currently not using an object, just trying to figure out how to do it. Using loops isn’t working. I tried different iterations, here’s where I’m at right now:
df_new = df.copy()
for country in median_data.Country:
country_data = median_data[median_data.Country == country].copy()
for col in median_data.columns[2:]:
df_new[col] = df_new[col].fillna(country_data[col])
The dataset, df:
The table with medians grouped by ‘Country’:
Result from the code above (clearly incorrect).
As an example, col4 for Country A should be [6, 6, 6], but the values I get are [6, 61, 6]:
Is there a way to separately calculate median values and impute them? Efficiency is not my primary concern, but an efficient solution would obviously be preferred.
>Solution :
Use DataFrame.update with GroupBy.transform with median:
df.update(df.groupby('Country').transform('median'), overwrite=False)
print (df)
Country col1 col2 col3 col4 col5 col6
0 A 10.0 20.0 4.0 6.0 50.0 30
1 A 2.0 1.0 5.0 6.0 34.0 35
2 A 13.0 212.0 3.0 6.0 42.0 37
3 B 120.0 230.0 53.0 61.0 63.0 23
4 B 22.0 115.0 15.0 61.0 4.0 15
5 B 71.0 22.0 12.0 61.0 33.5 31
6 C 105.0 120.0 14.0 22.0 520.0 3
7 C 26.0 11.0 15.0 128.0 34.0 3
8 C 13.0 65.5 13.0 234.0 277.0 10
9 D 101.0 220.0 654.0 143.0 634.0 123
10 D 32.0 21.0 61.0 24.0 584.0 32
11 D 11.0 72.0 23.0 83.5 534.0 30
Details:
print (df.groupby('Country').transform('median'))
col1 col2 col3 col4 col5 col6
0 10.0 20.0 4.0 6.0 42.0 35
1 10.0 20.0 4.0 6.0 42.0 35
2 10.0 20.0 4.0 6.0 42.0 35
3 71.0 115.0 15.0 61.0 33.5 23
4 71.0 115.0 15.0 61.0 33.5 23
5 71.0 115.0 15.0 61.0 33.5 23
6 26.0 65.5 14.0 128.0 277.0 3
7 26.0 65.5 14.0 128.0 277.0 3
8 26.0 65.5 14.0 128.0 277.0 3
9 32.0 72.0 61.0 83.5 584.0 32
10 32.0 72.0 61.0 83.5 584.0 32
11 32.0 72.0 61.0 83.5 584.0 32
Alternative solution with DataFrame.combine_first:
df1 = df.combine_first(df.groupby('Country').transform('median'))
print (df1)
Country col1 col2 col3 col4 col5 col6
0 A 10.0 20.0 4.0 6.0 50.0 30
1 A 2.0 1.0 5.0 6.0 34.0 35
2 A 13.0 212.0 3.0 6.0 42.0 37
3 B 120.0 230.0 53.0 61.0 63.0 23
4 B 22.0 115.0 15.0 61.0 4.0 15
5 B 71.0 22.0 12.0 61.0 33.5 31
6 C 105.0 120.0 14.0 22.0 520.0 3
7 C 26.0 11.0 15.0 128.0 34.0 3
8 C 13.0 65.5 13.0 234.0 277.0 10
9 D 101.0 220.0 654.0 143.0 634.0 123
10 D 32.0 21.0 61.0 24.0 584.0 32
11 D 11.0 72.0 23.0 83.5 534.0 30