groupby в pandas за 7 минут

groupby отвечает на вопросы вида «сколько заказов у каждого клиента» и «какая средняя выручка по городам». На собесе аналитика без него не обходится ни одна live-задача.

Split-apply-combine

Любой groupby состоит из трёх шагов:

1. Split: строки делятся на группы по значениям ключа.
2. Apply: к каждой группе применяется функция.
3. Combine: результаты склеиваются в новый объект.

import pandas as pd

df = pd.DataFrame({
    'city': ['Москва', 'Казань', 'Москва', 'Казань', 'Пермь'],
    'user': [1, 2, 1, 3, 4],
    'amount': [100, 50, 300, 70, 20],
})
df.groupby('city')['amount'].sum()
# city
# Казань     120
# Москва     400
# Пермь       20

Ключ ушёл в индекс, группы отсортированы по ключу (это делает sort=True по умолчанию). Если ключ нужен колонкой, добавьте as_index=False или вызовите .reset_index().

Три вида функций: agg, transform, filter

agg (aggregate) сворачивает группу в одно значение: одна строка на группу.

df.groupby('city').agg(
    orders=('amount', 'size'),
    revenue=('amount', 'sum'),
    users=('user', 'nunique'),
)

Это именованная агрегация: новое_имя=(колонка, функция). Колонки получаются плоскими, без MultiIndex, и сразу с понятными именами.

transform возвращает результат той же длины, что исходный DataFrame, и выровненный по его индексу. Это главный способ добавить к каждой строке групповую статистику:

df['city_total'] = df.groupby('city')['amount'].transform('sum')
df['share'] = df['amount'] / df['city_total']

Без transform пришлось бы сначала агрегировать, а потом делать merge обратно.

filter оставляет или выкидывает группы целиком:

df.groupby('city').filter(lambda g: len(g) >= 2)  # только города с 2+ заказами

size и count

• size() считает строки в группе, включая строки с NaN.
• count() считает непустые значения в каждой колонке.

Если в колонке есть пропуски, size и count разойдутся. Это любимый вопрос на собесе.

NaN в ключе

По умолчанию строки, у которых ключ группировки NaN, выпадают из результата (dropna=True). Если пропуск в ключе тоже значим (например, «город не указан»), передайте dropna=False.

apply: универсально, но дорого

apply вызывает Python-функцию на каждой группе, функция может вернуть что угодно: число, Series, DataFrame. Цена: на каждую группу Python-вызов, на миллионе групп это минуты. Порядок предпочтения:

1. встроенные строки 'sum', 'mean', 'max', 'nunique' (работают в C/Cython);
2. agg и transform со встроенными функциями;
3. apply только когда без него никак.

Например, топ-2 заказа на город можно сделать без apply:

df.sort_values('amount', ascending=False).groupby('city').head(2)

Несколько ключей и MultiIndex

df.groupby(['city', 'user'])['amount'].sum()

Результат получит MultiIndex (city, user). Дальше его можно развернуть в таблицу через .unstack() или вернуть в колонки через .reset_index().

Категориальные ключи и observed

Если ключ имеет тип category, pandas может показать и категории, которых нет в данных. Этим управляет параметр observed. В pandas 3.0 по умолчанию observed=True: показываются только встретившиеся категории. В старом коде с observed=False неожиданно появляются нулевые строки и раздувается MultiIndex.

Частые ошибки

• Забыли as_index=False и потом не могут сделать merge по колонке, которая уехала в индекс.
• Средняя от средних. Средний чек по городам, усреднённый ещё раз, не равен среднему чеку по всем заказам: у городов разное число заказов.
• apply вместо transform. apply с lambda g: g / g.sum() работает, но медленнее и с менее предсказуемым индексом.
• Путают count и size на колонке с пропусками.
• Цикл for name, group in df.groupby(...) для подсчёта сумм. Цикл по группам уместен для отладки или вывода, но не для вычислений.

Что спросят на собесе

• Чем agg, transform, apply и filter отличаются по форме результата?
• Как добавить к каждой строке долю от суммы её группы?
• Чем size отличается от count?
• Что происходит со строками, у которых ключ NaN?
• Как найти топ-N записей в каждой группе без apply?