badcasedaily1 23 сен в 07:09

Сложная агрегация в Pandas с MultiIndex

Простой

4 мин

Блог компании OTUSPython*Анализ и проектирование систем*

Обзор

Привет, Хабр!

Сегодня поговорим о Pandas MultiIndex — мощной штуке, которая спасла меня, когда пришлось агрегировать кучу многомерных данных. В отчётах за несколько лет с миллионами строк, обычный groupby() просто не справлялся, а MultiIndex сделал всё красиво и быстро. Разберём, как использовать его для сложных операций, не тратя лишнего времени.

Работа с MultiIndex

Создание MultiIndex

Создать MultiIndex в Pandas проще, чем может показаться. Есть несколько способов, и они зависят от того, с какими данными ты работаешь. Например, у тебя может быть DataFrame с несколькими столбцами, которые логически должны представлять вложенные индексы. В таком случае поможет метод set_index():

import pandas as pd

# пример данных с несколькими уровнями
data = {
    'country': ['USA', 'USA', 'Canada', 'Canada'],
    'city': ['New York', 'Los Angeles', 'Toronto', 'Vancouver'],
    'year': [2020, 2020, 2020, 2020],
    'population': [8.4, 4.0, 2.9, 2.5]
}

df = pd.DataFrame(data)

# Создание MultiIndex
df_multi = df.set_index(['country', 'city'])
print(df_multi)

На выходе у тебя будет DataFrame с двухуровневым индексом country и city. Теперь ты можешь выполнять операции на этих уровнях, как на обычных индексах, но с дополнительными возможностями.

Управление уровнями индексов

Окей, создали мы MultiIndex. Но что, если нужно добавить новый уровень или поменять местами уже существующие? Pandas предоставляет для этого несколько инструментов.

Добавление уровня индекса: метод set_index() можно использовать и для добавления индексов поверх уже существующих.

# Добавление индекса 'year' как еще одного уровня
df_multi = df_multi.set_index('year', append=True)
print(df_multi)

Переупорядочение уровней: для этого используем метод swaplevel(). Например, если нужно поменять местами страну и город:

df_reordered = df_multi.swaplevel('country', 'city')
print(df_reordered)

Удаление уровней индексов: это делается с помощью метода reset_index(), который "сбрасывает" один или несколько уровней индекса обратно в столбцы.

df_reset = df_multi.reset_index('city')
print(df_reset)

stack() и unstack(): работа с вложенными данными

Методы stack() и unstack() позволяют преобразовывать уровни индексов в столбцы и наоборот:

stack(): превращает уровни столбцов в уровни индексов. Мастхев для "поворота" данных в более компактный вид.

df_stacked = df_multi.stack()
print(df_stacked)

unstack(): обратная операция, превращающая индекс обратно в столбцы.

df_unstacked = df_multi.unstack()
print(df_unstacked)

Агрегация данных с MultiIndex

Теперь переходим к самому важному — как MultiIndex помогает выполнять сложные агрегации данных. Сложные многомерные группировки и агрегации можно реализовать на базе groupby() и agg(). Допустим, есть DataFrame с данными по годам и месяцам, и хотим агрегировать данные по нескольким уровням сразу.

Пример группировки по нескольким уровням:

# Допустим, у нас есть данные о продажах за несколько лет по странам и городам
data = {
    'country': ['USA', 'USA', 'Canada', 'Canada'],
    'city': ['New York', 'Los Angeles', 'Toronto', 'Vancouver'],
    'year': [2020, 2020, 2020, 2021],
    'sales': [100, 200, 150, 130]
}

df = pd.DataFrame(data)
df_multi = df.set_index(['country', 'city', 'year'])

# Группируем по странам и считаем общие продажи
df_grouped = df_multi.groupby(level='country').sum()
print(df_grouped)

Этот код покажет нам суммарные продажи по странам за все годы. Но можно пойти еще дальше и агрегировать данные по уровням внутри уровней. Например, используя agg(), можно комбинировать разные типы агрегации:

df_aggregated = df_multi.groupby(level='country').agg({'sales': ['sum', 'mean']})
print(df_aggregated)

Здесь считаем одновременно и сумму, и среднее по продажам для каждого уровня.

Работа с срезами данных по уровням MultiIndex

Когда есть MultiIndex, стандартные методы доступа к данным типа .loc[] становятся еще более мощными. Например, можно извлечь данные по одному уровню или по комбинации уровней.

Доступ к данным через .loc[]:

# Доступ к данным для определенного города
print(df_multi.loc[('USA', 'New York')])

Использование .xs() для выборки данных по конкретному уровню:

# Получение данных по конкретному году
print(df_multi.xs(2020, level='year'))

Прочие полезные фичи

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

Проверка наличия MultiIndex: иногда нужно быстро проверить, есть ли MultiIndex в DataFrame.

print(df_multi.index.is_multi)

Получение имен уровней:

print(df_multi.index.names)

Распаковка уровней индексов в столбцы для упрощения работы:

df_flat = df_multi.reset_index()
print(df_flat)

Совмещение нескольких индексов в один уровень: когда нужно скомбинировать несколько индексов в один, можно использовать pd.MultiIndex.from_tuples().

new_index = pd.MultiIndex.from_tuples([('USA', 2020), ('Canada', 2021)], names=['country', 'year'])

Вот так Pandas MultiIndex превращает сложные многомерные данные в легко управляемые структуры, где агрегация и манипуляция уровнями выполняются в пару строк. Если ваши данные сложнее обычных таблиц, MultiIndex — это ваш лучший друг.

Освоить мощные навыки анализа данных — анализ требований + статистика + BI и получить востребованную профессию можно на курсе «Аналитик данных».

Теги:

Хабы: