I have two dataframes df_country_cluster
and df_countries
, with the following structure (in order):
cluster_id
country_id
1 4
2 4
... ...
col1 col2 col3 col4
country_id year_id
1 2015 0.1 0.2 0.3 0.1
1 2016 0.4 NaN 0.1 0.8
1 2017 0.7 0.2 0.6 NaN
1 2018 0.9 0.4 0.7 0.2
2 2015 0.5 0.6 NaN 0.3
2 2016 0.3 0.7 0.2 0.5
2 2017 0.2 0.9 0.3 0.5
2 2018 0.1 0.2 0.1 0.9
... ... ... ... ... ...
My goal is to fill the NaN values with the average non-NaN values of each cluster and year. This means that, for example, the NaN in country_id
1, year_id
2016, col2
should be filled with the average of the valid values of col2
for 2016 and all the countries with cluster_id
4 (in this case).
In the above example, we would get the averages of cluster_id
4 in this way:
col1 col2 col3 col4
cluster_id year_id
4 2015 0.3 0.4 *0.3 0.2
4 2016 0.4 *0.7 0.2 0.6
4 2017 0.4 0.6 0.4 *0.5
4 2018 0.5 0.3 0.4 0.6
Therefore, the NaN each column would be filled by the values with a *.
I have tried to create a new DataFrame
with groupby().mean()
, and then use .fillna
, but without success. Other SO questions like this only discuss the single-index problem.
Here is my approach:
cols = ['col1','col2','col3','col4']
original_index = df_countries.index
df_countries = df_countries.join(df_country_cluster,on='country_id')
df_countries = df_countries.reset_index().set_index(['cluster_id','year_id'])
avg_cluster = df_countries.groupby(['cluster_id','year_id']).mean()
avg_cluster = avg_cluster[cols]
for col in cols:
df_countries[col].fillna(avg_cluster[col],inplace=True)
df_countries.reset_index().set_index(original_index)