Панды: заменить NaN со средним на многоуровневый индекс

Question

У меня есть два кадра данных df_country_cluster и df_countries со следующей структурой (по порядку):

            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 
...         ...      ...    ...    ...    ...

Моя цель - заполнить значения NaN средними значениями, отличными от NaN, для каждого кластера и года. Это означает, что, например, значение NaN в country_id 1, year_id 2016, col2 должно быть заполнено средним значением действительных значений col2 для 2016 года и всех стран с cluster_id 4 (в это дело).

В приведенном выше примере мы получили бы средние значения cluster_id 4 следующим образом:

                     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

Следовательно, NaN каждого столбца будет заполнен значениями с *.

Я попытался создать новый DataFrame с groupby().mean(), а затем использовать .fillna, но безуспешно. Другие вопросы SO , подобные этому , обсуждают только проблему с одним индексом.

Вот мой подход:

    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)

Answer 1

Понял.

df_countries = df_countries.reset_index().set_index(original_index)

Забыл оставить ответ с правильным индексом ... С этим изменением все работает. Однако, если у кого-то есть более питонический способ сделать это, пожалуйста, добавьте свой ответ!

Answer 2

Я не уверен, правильно ли я вас понял.Но ваш подход выглядит хорошо.Итак, мы начинаем с

df = df_countries.join(df_country_cluster, on='country_id')
df = df.reset_index().set_index(['cluster_id','year_id'])

Вы говорите, что значение заполнения является средним для каждой группы:

s = df[['col1', 'col2', 'col3', 'col4']].mean(axis=1)

Мы можем сделать все вместе, если мы транспонируем DataFrame

df = df.T.fillna(value=s).T

Наконец мы отбрасываем то, что нам не нужно

df = df.reset_index().drop(columns='cluster_id').set_index('country_id', 'year_id')