Как лексикографически отсортировать группы в pandas

Пошаговый способ лексикографической сортировки групп в pandas: сортировка внутри групп, агрегация в кортежи, сортировка по ним и explode. Пример кода и нюансы.

Когда нужно отсортировать группы в pandas не по одному агрегату, а по самим значениям по возрастанию с корректным разрешением равенств, простого минимума на группу недостаточно. Требуется лексикографическая сортировка между группами: сравниваем наименьшее значение; если оно совпадает — берём следующее, и так далее, пока порядок не определится.

Пример набора данных

Рассмотрим небольшой датафрейм. Задача — упорядочить группы по значениям arrive, используя последующие значения как тайбрейкеры, а затем вывести строки в этом порядке.

import pandas as pd
import numpy as np
tbl = pd.DataFrame({
    "group_id": [5, 1, 9, 9, 5, 7, 7, 7, 9, 1, 5],
    "arrive":   [227, 60, 60, 88, 55, 55, 276, 46, 46, 35, 35]
})

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

Отсортировать внутри каждой группы несложно, но здесь требуется сортировка между группами по последовательности их значений. Опираться только на минимум группы нельзя: если у двух групп одинаковое первое значение, но далее они расходятся, порядок будет неверным. Попытки комбинировать groupby с transform и nth тоже не срабатывают: передача "nth" в transform приводит к ошибке, потому что для одной группы он может вернуть ни одного или сразу несколько значений.

Один из способов — добавить вспомогательные столбцы, где будут зафиксированы первое, второе, третье значения в группе, и сортировать по ним. Это работает, но плохо масштабируется, если в группах много элементов.

Императивный обходной путь (работает, но громоздок)

Ниже показан подробный подход: создаются столбцы для каждой позиции и выполняется сортировка по ним. Логика верная, но поддерживать N позиционных столбцов при больших группах неудобно.

# сортируем по значению, чтобы задать порядок внутри каждой группы
wrk = tbl.sort_values("arrive").copy()
wrk["rank_in_group"] = wrk.groupby("group_id")["arrive"].cumcount()
# фиксируем первые три позиции
wrk["a1"] = wrk["a2"] = wrk["a3"] = np.nan
wrk.loc[wrk["rank_in_group"] == 0, "a1"] = wrk.loc[wrk["rank_in_group"] == 0, "arrive"]
wrk.loc[wrk["rank_in_group"] == 1, "a2"] = wrk.loc[wrk["rank_in_group"] == 1, "arrive"]
wrk.loc[wrk["rank_in_group"] == 2, "a3"] = wrk.loc[wrk["rank_in_group"] == 2, "arrive"]
# распространяем позиционные значения на все строки той же группы
wrk[["a1", "a2", "a3"]] = (
    wrk.groupby("group_id")[ ["a1", "a2", "a3"] ].transform("max")
)
# для коротких групп заполняем пропуски предыдущими позициями
wrk["a2"] = wrk["a2"].fillna(wrk["a1"])
wrk["a3"] = wrk["a3"].fillna(wrk["a2"])
# финальная сортировка по позиционным ключам и удаление вспомогательных столбцов
out_verbose = wrk.sort_values(["a1", "a2", "a3", "group_id"]) \
               .drop(columns=["a1", "a2", "a3", "rank_in_group"]) 

Ключевая идея

Суть задачи — лексикографическое сравнение отсортированных значений внутри группы. Если преобразовать значения каждой группы в один сравнимый объект, сохраняющий порядок, можно отсортировать такие объекты, а затем развернуть их обратно в строки. Кортежи подходят идеально, и pandas умеет агрегировать столбцы в кортежи напрямую.

Компактное решение: агрегировать в кортежи, отсортировать, развернуть

Подход ниже сначала сортирует значения внутри групп, превращает каждую группу в упорядоченный кортеж, сортирует по этому кортежу, а затем разворачивает его обратно в строки. Так мы получаем нужный порядок между группами, опираясь на стандартные средства.

lexi = (
    tbl.sort_values("arrive")
       .groupby("group_id", as_index=False)
       .agg(tuple)
)
# lexi
#    group_id           arrive
# 0         1         (35, 60)
# 1         5    (35, 55, 227)
# 2         7    (46, 55, 276)
# 3         9      (46, 60, 88)
result = (
    tbl.sort_values("arrive")
       .groupby("group_id", as_index=False)
       .agg(tuple)
       .sort_values("arrive")
       .explode("arrive")
)
# результат
#    group_id arrive
# 1         5     35
# 1         5     55
# 1         5    227
# 0         1     35
# 0         1     60
# 2         7     46
# 2         7     55
# 2         7    276
# 3         9     46
# 3         9     60
# 3         9     88

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

Сортировка между группами — иная задача, чем упорядочивание внутри группы. Если рассматривать отсортированные значения группы как единый упорядоченный объект, мы получаем лексикографическое сравнение и избегаем хрупких конструкций с множеством позиционных столбцов. Важно помнить и о производительности: хотя решение элегантно и лаконично, работа с groupby и объектами Python может быть медленнее векторных операций; даже сортировка кортежей нередко заметно медленнее сортировки числовых серий в простых случаях. Понимание как корректности, так и возможных издержек помогает выбрать подход под вашу задачу.

Итоги

Если вам нужна детерминированная, учитывающая равенства сортировка между группами в pandas, сначала отсортируйте значения в группах, затем агрегируйте их в кортежи, отсортируйте по этим кортежам и разверните обратно в строки. Такой подход точно отражает семантику лексикографического порядка при минимуме кода. При больших данных и строгих требованиях к задержкам учитывайте накладные расходы объектного уровня в Python и подумайте, подходит ли более векторизованный путь под ваши ограничения.