Условие

Дан DataFrame events с колонками user_id, event_date, event_type. Постройте retention curve:

Когорта — неделя первого события.
По оси X — week_offset (0, 1, 2, ...).
По оси Y — % пользователей когорты, активных в эту неделю.

Верните DataFrame cohort_week × week_offset → pct_retained.

Решение

import pandas as pd
import numpy as np
 
def retention(events: pd.DataFrame,
              max_weeks: int = 12) -> pd.DataFrame:
    df = events.copy()
    df['event_date'] = pd.to_datetime(df['event_date'])
    # начало недели (понедельник)
    df['event_week'] = df['event_date'].dt.to_period('W-MON').dt.start_time
 
    # cohort = неделя первого события
    cohort = df.groupby('user_id')['event_week'].min().rename('cohort_week')
    df = df.merge(cohort, on='user_id')
 
    df['week_offset'] = ((df['event_week'] - df['cohort_week']).dt.days // 7).astype(int)
    df = df[df['week_offset'].between(0, max_weeks)]
 
    # уникальные user_id в каждой ячейке
    active = (df.groupby(['cohort_week', 'week_offset'])['user_id']
                .nunique()
                .reset_index(name='active'))
 
    sizes = (active.query('week_offset == 0')[['cohort_week', 'active']]
                   .rename(columns={'active': 'cohort_size'}))
 
    out = active.merge(sizes, on='cohort_week')
    out['pct_retained'] = (out['active'] / out['cohort_size']) * 100
 
    pivot = (out.pivot(index='cohort_week',
                       columns='week_offset',
                       values='pct_retained')
                 .round(1))
    pivot.index = pivot.index.strftime('%Y-%m-%d')
    return pivot

Пример

events = pd.DataFrame({
    'user_id':    [1,1,1,2,2,3,3,4,4,4],
    'event_date': ['2025-01-06','2025-01-08','2025-01-20',
                   '2025-01-06','2025-01-13',
                   '2025-01-13','2025-02-03',
                   '2025-01-13','2025-01-20','2025-02-10'],
    'event_type': ['signup']*10,
})
print(retention(events, max_weeks=4))

Ожидаемый вывод (пример):

week_offset    0     1     2     3     4
cohort_week
2025-01-06   100.0  50.0   0.0  50.0   0.0
2025-01-13   100.0  50.0   0.0   0.0  50.0

Вариация: retention по конкретному event_type

def retention_by_event(events, target_event='purchase', max_weeks=12):
    df = events[events['event_type'] == target_event].copy()
    # ... остальной код тот же

Визуализация heatmap

import matplotlib.pyplot as plt
import seaborn as sns
 
ret = retention(events, max_weeks=12)
plt.figure(figsize=(12, 6))
sns.heatmap(ret, annot=True, fmt='.0f', cmap='YlGn', vmin=0, vmax=100)
plt.title('Cohort retention, %')
plt.show()

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

dt.to_period('W') в pandas по умолчанию использует W-SUN (неделя заканчивается в воскресенье). Для ISO-понедельника — W-MON + .dt.start_time.
merge вместо map: для cohort_week часто пишут df['cohort'] = df['user_id'].map(cohort_dict) — работает, но при больших данных merge быстрее.
Свежие когорты «не созрели»: для cohort_week = 2025-05-19 неделя 4 ещё не наступила → ячейка должна быть NaN, не 0. После pivot — естественно NaN.
Pivot для динамических колонок: если когорт > 50 и недель > 24 — таблица будет огромной; ограничивать max_weeks.
nunique дорого на сотнях миллионов строк → заранее свернуть до user_id × event_week.
Часовые пояса: pd.to_datetime теряет TZ если строка без offset. Для timestamptz из БД — utc=True, потом dt.tz_convert(...).
week_offset = (week_diff).days // 7 работает только если обе даты — начала недели. Иначе будет off-by-one.

Эталонный ответ

Функция retention(events, max_weeks): (1) приводим event_date к началу недели; (2) cohort = groupby('user_id').event_week.min(); (3) week_offset = (event_week - cohort_week).days // 7; (4) nunique user_id в каждой ячейке; (5) делим на размер когорты (week_offset=0); (6) pivot → таблица.

Подводные камни: W-MON для ISO-недель; свежие когорты «не созрели» (NaN, не 0); часовые пояса; nunique дорого на больших данных.

Karpov ДЗ: Retention curve на pandas