Python: Get key value pairs within a group in pandas DataFrame

user_7122272 0
Viewed 147

I have the following data frame:

import pandas as pd

df = pd.DataFrame([
    {'car_id': 123, 'country_code': 'CZ', 'grade': 5.0},
    {'car_id': 123, 'country_code': 'SK', 'grade': 1.0},
    {'car_id': 123, 'country_code': 'PL', 'grade': 4.0},
    {'car_id': 234, 'country_code': 'CZ', 'grade': 4.0},
    {'car_id': 234, 'country_code': 'SK', 'grade': 2.0},
    {'car_id': 234, 'country_code': 'PL', 'grade': 3.0},
    {'car_id': 345, 'country_code': 'CZ', 'grade': 2.0},
    {'car_id': 345, 'country_code': 'SK', 'grade': 5.0},
    {'car_id': 345, 'country_code': 'PL', 'grade': 1.0},
    {'car_id': 456, 'country_code': 'CZ', 'grade': None},
    {'car_id': 456, 'country_code': 'SK', 'grade': None},
    {'car_id': 456, 'country_code': 'PL', 'grade': None}
])

Now I want to group data by car_id and get two columns:

  • minimal grade,
  • country code of the minimal grade.

So far, I have the following code:

>>> (
...     df
...     .groupby('car_id')
...     .apply(lambda x: pd.Series({
...         'min_grade': x['grade'].min(),
...         'min_grade_country': x.loc[x.grade == x.grade.min(), 'country_code'],
...     }))
...     .reset_index()
... )
 car_id  min_grade   min_grade_country
0   123        1.0   1 SK Name: country_code, dtype: object
1   234        2.0   4 SK Name: country_code, dtype: object
2   345        1.0   8 PL Name: country_code, dtype: object
3   456        NaN   Series([], Name: country_code, dtype: object)

As you can see, I cannot extract the country code with the minimal grade. And also, I am not sure if there is any more elegant pandas way how to get it - I mean without using .apply() method in combination with lambda function. Can you help me with that, please?

pandas group-by python
1 Answers

Use GroupBy.agg with aggregate by min and DataFrameGroupBy.idxmin for indices from country_code converted to index first:

df1 = (
    df
    .set_index('country_code')
    .groupby('car_id')
    .agg(
        min_grade=('grade', 'min'),
        min_grade_country=('grade', 'idxmin')
    )
    .reset_index()
)

print (df1)
   car_id  min_grade min_grade_country
0     123        1.0                SK
1     234        2.0                SK
2     345        1.0                PL
3     456        NaN               NaN
user_29010020
Related