pandas фильтр строк на основе поиска максимума в столбце

Question

Задача такая. Есть датафрейм:

df123 = pd.DataFrame({"A": ["1", "1", "1", "2", "2"],
                      "B": ["monday", "tuesday", "wednesday", "monday", "tuesday"],
                      "C": ["1", "2", "3", "2", "2"]})

Нужно для каждого уникального значения колонки А вывести строку с максимальным значением колонки С, при этом колонка В также должна показываться. Если одинаковых максимумов несколько, то нужно показать все сроки с одинаковыми максимумами (они будут отличаться по столбцу В) Т.е. на выходе должно быть:

Мне в голову приходит только

    df123.groupby('A', as_index = False)\
     .agg({'C' : 'max'})

Но тогда теряю колонку В. Чувствую, где-то на поверхности плавает решение, не могу сообразить.

Answer 1

df123.groupby(['A']).apply(lambda t: t[t.C==t.C.max()])

Результат:

     A          B  C
A                   
1 2  1  wednesday  3
2 3  2     monday  2
  4  2    tuesday  2

P.S. Из за жалоб на медленность работы провел эксперимент и с большим удивлением обнаружил, что результат можно немного ускорить за счет весьма незначительной переделки:

%timeit df123.groupby(['A']).apply(lambda t: t[t.C==t.C.max()])

1.85 ms ± 9.09 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit df123.groupby(['A']).apply(lambda t: t[t.C.values==t.C.values.max()])

1.62 ms ± 14.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Answer 2

Странно, что никто не вспомнил про merge:

pd.merge(df123, df123.groupby('A').agg({'C': 'max'}).reset_index(), on=('A','C'))

Вывод:

    A   B           C 
0   1   wednesday   3
1   2   monday      2
2   2   tuesday     2

Работает медленнее, чем решение с apply, но возможно в каких-то случаях это решение будет более универсальным.

БЛОГ НА HUSL

pandas фильтр строк на основе поиска максимума в столбце

Ответы (2 шт):