groupby и transform: групповые признаки на уровне строк и фильтрация групп

Добавление группового показателя как нового столбца

Процедура укладывается в два шага: вызвать transform и присвоить результат новому столбцу.

import pandas as pd

df = pd.DataFrame({
    'dept': ['HR', 'IT', 'HR', 'IT', 'HR'],
    'salary': [60000, 120000, 70000, 110000, 80000]
})

df['dept_mean'] = df.groupby('dept')['salary'].transform('mean')
df['dept_sum']  = df.groupby('dept')['salary'].transform('sum')

print(df)

Вывод:

  dept  salary   dept_mean  dept_sum
0   HR   60000   70000.0    210000
1   IT  120000  115000.0    230000
2   HR   70000   70000.0    210000
3   IT  110000  115000.0    230000
4   HR   80000   70000.0    210000

Что здесь происходит. Pandas разбивает строки по значению dept, считает среднее и сумму внутри каждой группы, а затем возвращает значение обратно в каждую строку по её исходной позиции. HR получает среднее (60000 + 70000 + 80000) / 3 = 70000, IT — (120000 + 110000) / 2 = 115000. Длина результата — ровно 5, как у исходного df.

В строку присваивания df['dept_mean'] = ... результат transform ложится без каких-либо дополнительных операций — индексы совпадают, pandas выравнивает значения автоматически.

Аргумент func в transform — строковый псевдоним агрегата: 'mean', 'sum', 'max', 'min', 'std' и другие стандартные агрегатные имена, которые принимает pandas. Передавать туда NumPy-функцию или лямбду тоже можно, но строковые имена работают быстрее, потому что pandas использует оптимизированные пути вычисления.