Python Dataframe 为基于组最大值的列设置值

Question

I have the following dataframe

   id   Area Country
0  11  34.45  Norway
1  12  30.25      UK
2  13  16.70    Iran
3  11  35.45  Sweden
4  13  20.22    Iraq
5  15  35.12     USA

dfObj['BigCountry'] = ''
dfObj['SmallCountry'] = ''

Based on the area I want to classify the country either big or small. I was trying to groupby id and then based on max(area) within the group I want to set small/big country

The output should be

   id  BigCountry  SmallCountry   
0  11  Sweden         Norway
1  12  UK             UK           
2  13  Iraq           Iran
5  15  USA            USA 

Answer 1

One way is to use set_index then groupby and agg with idxmax and idxmin :

df.set_index('Country').groupby('id')['Area'].agg(['idxmax','idxmin'])\
  .rename(columns = {'idxmax':'BigCountry', 'idxmin':'SmallCountry'})

Output:

   BigCountry SmallCountry
id                        
11     Sweden       Norway
12         UK           UK
13       Iraq         Iran
15        USA          USA