Группировка, применение функции и вставка нового столбца с соответствующими значениями

Question

Я пытаюсь вычислить средневзвешенное значение для определенных групп пользователей, но, похоже, не могу взять вывод и сформировать новый столбец с ним.

Я пробовал группировать и применять функцию, но могуне похоже, чтобы преобразовать его обратно в новый столбец

 df = pd.DataFrame('cohort_size': [10,15,20,25,20,15],
'retention' : [0.3,0.2,0.4,0.5,0.6,0.5],
'nth_interval' : [1,1,1,7,7,7])

df.groupby('nth_interval').apply(lambda x: np.average(x['retention'], weights=x['cohort_size']))

Я получаю выходные данные в виде серии, но я действительно хочу взять эти индивидуальные взвешенные средние и вернуть их в виде нового столбца вфрейм данных, где каждому значению столбца соответствует nth_interval

Answer 1

Вы можете сделать слияние:

df.merge(df.groupby('nth_interval')
           .apply(lambda x: np.average(x['retention'], weights=x['cohort_size']))
           .to_frame(name='average'),
         on='nth_interval')

или карту:

s = df.groupby('nth_interval')
               .apply(lambda x: np.average(x['retention'], weights=x['cohort_size']))

df['average'] = df['nth_interval'].map(s)

вывод:

   cohort_size  retention  nth_interval   average
0           10        0.3             1  0.311111
1           15        0.2             1  0.311111
2           20        0.4             1  0.311111
3           25        0.5             7  0.533333
4           20        0.6             7  0.533333
5           15        0.5             7  0.533333

Или вы также можете вычислить это вручную:

df['average'] = (df['retention'].mul(df['cohort_size'])
                      .groupby(df['nth_interval'])
                      .transform('sum')
                      .div(df['cohort_size'].groupby(df['nth_interval'])
                                            .transform('sum'))
                 )

Answer 2

IIUC, вы ищете это?

df.merge(df.groupby('nth_interval').apply(lambda x: np.average(x['retention'],
                                                   weights=x['cohort_size']))
                                   .rename('wgted_avg').reset_index(), 
         on='nth_interval')

Вывод:

   cohort_size  retention  nth_interval  wgted_avg
0           10        0.3             1   0.311111
1           15        0.2             1   0.311111
2           20        0.4             1   0.311111
3           25        0.5             7   0.533333
4           20        0.6             7   0.533333
5           15        0.5             7   0.533333