Постройте когортную retention-матрицу в pandas с помощью pivot table?

Короткий ответ

Нужно для каждой пары (пользователь, событие) вычислить номер периода когорты, сгруппировать по (cohort_month, period_number), посчитать число уникальных пользователей, а затем развернуть результат через pivot_table в матрицу "когорта × период" и нормировать на размер когорты.

Как это работает подробнее

после группировки получаем длинную (long-format) таблицу с тремя колонками (cohort_month, period_number, active_users); pivot_table разворачивает её в широкий формат, где строки — когорты, столбцы — номер периода, значения — число активных или retention rate.

Пример кода

cohort_retention.pyКопировать

import pandas as pd
df = pd.DataFrame({
    'user_id': [1, 1, 2, 2, 3],
    'event_date': pd.to_datetime(['2024-01-05', '2024-02-10', '2024-01-20', '2024-03-01', '2024-02-02']),
})
df['cohort_month'] = df.groupby('user_id')['event_date'].transform('min').dt.to_period('M')
df['event_month'] = df['event_date'].dt.to_period('M')
df['period_number'] = (df['event_month'] - df['cohort_month']).apply(lambda x: x.n)
cohort_data = df.groupby(['cohort_month', 'period_number'])['user_id'].nunique().reset_index()
cohort_matrix = cohort_data.pivot_table(index='cohort_month', columns='period_number', values='user_id', aggfunc='sum')
retention_matrix = cohort_matrix.div(cohort_matrix.iloc[:, 0], axis=0).round(2)
print(retention_matrix)

Что использовать на практике

• Практика. в pandas 2.x при работе с Period/Grouper для месячных когорт стоит явно приводить типы (dt.to_period('M')), а после merge/groupby проверять df.shape, чтобы не упустить незаметное дублирование строк; для нормировки на размер когорты удобно делить матрицу на первый столбец (.div(matrix.iloc[:, 0], axis=0)).

Подводные камни

• Частые ошибки. забывают нормировать на размер когорты и путают абсолютные числа активных пользователей с retention rate (долей); при pivot_table по умолчанию агрегирующая функция — mean, а не sum/count, что даёт неверный результат, если явно не указать aggfunc.