Группировка (groupby) в pandas c сохранением количества строк, генерация агрегатных фич.

Пусть у нас есть датафрейм и нам надо сгруппировать его по колонке «A»

import pandas as pd
import numpy as np


df = pd.DataFrame()
df['A'] = [1, 1, 1, 1, 1, 2, 2, 3, 3, 3]
df['v'] = [3, 5, 7, 10, 20, 11, 1, 3, 8, 10]

print(df.head(10))
# распечатаем датафрейм для наглядности
   A   v
0  1   3
1  1   5
2  1   7
3  1  10
4  1  20
5  2  11
6  2   1
7  3   3
8  3   8
9  3  10

Для группировка выберем функцию среднего и после применения операции группировки в DataFrame остается всего 2 колонки. Пример:

res = df.groupby(by=['A'], as_index=False)['v'].mean()
print(res.head())
# результат
   A  v
0  1  9
1  2  6
2  3  7

Как сделать так, чтобы просто добавить результат агрегации в новую колонку? Для этого воспользуемся функцией transform

df['col_mean'] = df.groupby(by=['A'], as_index=False)['v'].transform(lambda s: np.mean(s.values))
print(df.head(10))
# получаем результат агрегации в новой колонке
   A   v  col_mean
0  1   3       9.0
1  1   5       9.0
2  1   7       9.0
3  1  10       9.0
4  1  20       9.0
5  2  11       6.0
6  2   1       6.0
7  3   3       7.0
8  3   8       7.0
9  3  10       7.0

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

import pandas as pd
import numpy as np
import scipy
from scipy import stats


def make_agg(df_in, group_col_name, stat_col_name):
    df = df_in.copy()
    # то же самое что и count, количество
    df['col_count'] = df.groupby(by=[group_col_name], as_index=False)[stat_col_name].transform(lambda s: s.values.shape[0])
    # среднее значение по сгруппированному столбцу
    df['col_mean'] = df.groupby(by=[group_col_name], as_index=False)[stat_col_name].transform(lambda s: np.mean(s.values))
    # медиана по сгруппированному столбцу
    df['col_median'] = df.groupby(by=[group_col_name], as_index=False)[stat_col_name].transform(lambda s: np.median(s.values))
    # сумма по сгруппированному столбцу
    df['col_sum'] = df.groupby(by=[group_col_name], as_index=False)[stat_col_name].transform(lambda s: np.sum(s.values))
    # минимум и максимум по сгруппированному столбцу
    df['col_max'] = df.groupby(by=[group_col_name], as_index=False)[stat_col_name].transform(lambda s: np.max(s.values))
    df['col_min'] = df.groupby(by=[group_col_name], as_index=False)[stat_col_name].transform(lambda s: np.min(s.values))
    # стандартное отклонение по сгруппированному столбцу
    df['col_std'] = df.groupby(by=[group_col_name], as_index=False)[stat_col_name].transform(lambda s: np.std(s.values))
    # квантили в 20 и 80 % по сгруппированному столбцу
    df['col_q20'] = df.groupby(by=[group_col_name], as_index=False)[stat_col_name].transform(lambda s: np.quantile(s.values, q=0.2))
    df['col_q80'] = df.groupby(by=[group_col_name], as_index=False)[stat_col_name].transform(lambda s: np.quantile(s.values, q=0.8))
    # skew
    df['col_skew'] = df.groupby(by=[group_col_name], as_index=False)[stat_col_name].transform(lambda s: scipy.stats.skew(s))

    return df

res = make_agg(df, 'A', 'v')
print(res.head(10))

Как видите в transform можно написать любую свою lambda функцию. Только помните, что переменная придет типом Series, и для того чтобы случайно не запутаться в индексах, рекомендую добавлять values, тем самым переходя к стандартному numpy массиву. Также считать агрегатные статистики лучше сразу по объединенному train+test датасету, а не по отдельности для train и test. Так как объединенный датасет будет более показательным.

Добавить комментарий

Ваш e-mail не будет опубликован. Обязательные поля помечены *

Можно использовать следующие HTML-теги и атрибуты: <a href="" title=""> <abbr title=""> <acronym title=""> <b> <blockquote cite=""> <cite> <code> <del datetime=""> <em> <i> <q cite=""> <strike> <strong>