Pandas groupby — получить выходное значение на основе максимального значения другого столбца

#python #pandas #sklearn-pandas

Вопрос:

У меня есть следующий фрейм данных:

 df = pd.DataFrame({'Animal': ['Falcon', 'Falcon',
                              'Parrot', 'Parrot'],
                   'Habitat':['Jungle', 'Jungle',
                              'Sky', 'Sky'],
                   'Tmp':['A', 'B', 'C', 'D'],
                   'Max Speed': [380., 370., 24., 26.]})

>>> df
   Animal Habitat Tmp  Max Speed
0  Falcon  Jungle   A      380.0
1  Falcon  Jungle   B      370.0
2  Parrot     Sky   C       24.0
3  Parrot     Sky   D       26.0
 

Я пытаюсь добавить дополнительный столбец «Вывод», который будет принимать значение из «Tmp» на основе максимального значения столбца «Максимальная скорость» в группе, состоящей из столбцов «Животное» и «Среда обитания».

Желаемый результат:

    Animal Habitat Tmp  Max Speed Output
0  Falcon  Jungle   A      380.0      A
1  Falcon  Jungle   B      370.0      A
2  Parrot     Sky   C       24.0      D
3  Parrot     Sky   D       26.0      D
 

Это можно сделать с помощью a groupby , а затем объединить его с исходным набором данных. Но есть ли более эффективный способ сделать это? Может быть, использовать transform или что-то еще?

Ответ №1:

Вы можете определить функцию, принимающую pd.dataframe в качестве аргумента:

 import pandas as pd
import numpy as np

def fmax(df_):
    df_['Output'] = df_.sort_values(['Max Speed']).tail(1)['Tmp'].squeeze()
    return df_
 

Пожалуйста, обратите внимание на использование pandas.DataFrame.squeeze функции для возврата скалярного значения.
Затем просто apply используйте функцию выше, используя groupby :

 df.groupby(['Animal','Habitat']).apply(fmax)
 

В результате получается:

    Animal Habitat Tmp  Max Speed Output
0  Falcon  Jungle   A      380.0      A
1  Falcon  Jungle   B      370.0      A
2  Parrot     Sky   C       24.0      D
3  Parrot     Sky   D       26.0      D