Вычитание в pandas по составному индексу: set_index и sub

Показываем, как корректно вычитать DataFrame в pandas по составному ключу: выравнивание через set_index и sub с fill_value=0. Пример кода и ожидаемый результат.

Когда нужно корректировать агрегированные показатели по нескольким ключам в pandas, наивное вычитание быстро дает сбои. Если у двух DataFrame разный индекс, DataFrame.sub выравнивает данные по меткам индекса и названиям столбцов, а не по составному ключу вроде School, District, Program, Grade и Month. Итогом становятся неверные вычитания между строками или неожиданные NaN. Решение простое и надежное — выровнять данные по индексу из нескольких столбцов перед вычитанием.

Постановка задачи

Представьте: в одном DataFrame — общая посещаемость, в другом — величины, которые нужно вычесть для конкретных сочетаний ключей. Цель — вычитать только там, где совпадают School, District, Program, Grade и Month, оставляя остальное без изменений.

import pandas as pd

base_df = pd.DataFrame({
    'School':   [123, 123, 321, 321],
    'District': [456, 456, 654, 456],
    'Program':  ['A',  'B',  'A',  'A'],
    'Grade':    ['9-12','9-12','9-12','7-8'],
    'Month':    [10,   10,   10,   10],
    'Count':    [100,  95,   23,   40]
})

subtract_df = pd.DataFrame({
    'School':   [123, 321],
    'District': [456, 654],
    'Program':  ['A',  'A'],
    'Grade':    ['9-12','9-12'],
    'Month':    [10,   10],
    'Count':    [10,   8]
})

Если попытаться вычесть таблицы напрямую, pandas будет сопоставлять строки по умолчанию по целочисленному индексу и названиям столбцов, а не по предметным ключам. Это не то, что нам нужно.

Почему наивный подход не работает

Арифметика между DataFrame учитывает метки. При отсутствии общего индекса pandas сопоставляет строки по значениям индекса (0, 1, 2, …) и столбцы по именам. В этой задаче строки должны сопоставляться по составному ключу из пяти столбцов. Без такого выравнивания операция может вычитать несвязанные строки или давать NaN там, где индексы не совпадают. Для многих сценариев это корректное поведение, но не для сопоставления записей по нескольким ключам.

Решение: выровнять по составному индексу и только потом вычитать

Надежный способ: временно сделать ключевые столбцы индексом в обоих DataFrame, выполнить вычитание с параметром fill_value=0, чтобы отсутствующие соответствия считались нулями, а затем вернуть столбцы на место.

key_cols = ['School', 'District', 'Program', 'Grade', 'Month']

result_df = (
    base_df.set_index(key_cols)
           .sub(subtract_df.set_index(key_cols), fill_value=0)
           .reset_index()
)

Так строки сопоставляются строго по School, District, Program, Grade и Month. Любой ключ, присутствующий лишь в одном из DataFrame, корректно обрабатывается благодаря fill_value=0: там, где нет пары, вычитание не выполняется.

Ожидаемый результат

Результат сохраняет все исходные строки и вычитает только там, где ключи совпадают:

   School  District Program Grade  Month  Count
0     123       456       A  9-12     10   90.0
1     123       456       B  9-12     10   95.0
2     321       456       A   7-8     10   40.0
3     321       654       A  9-12     10   15.0

Почему это важно

Операции, кажущиеся простыми на первый взгляд, становятся хрупкими, когда задействовано несколько измерений. В pandas корректность зависит от выравнивания. Явно задавая и используя составной индекс, вы получаете предсказуемое, повторяемое поведение и избегаете построчной логики, которая плохо масштабируется. Этот прием легко переносится на корректировки, сверки и любые вычисления, где уникальность определяется набором столбцов.

Выводы

Перед арифметикой между DataFrame явно задайте критерии сопоставления — вынесите нужные столбцы в индекс. Используйте sub с fill_value=0, чтобы аккуратно обработать непересекающиеся ключи, затем сбросьте индекс и вернитесь к привычной форме столбцов. Подход лаконичен, сохраняет исходную логику группировки и гарантирует, что изменятся только нужные записи.