Search code examples
pythonpandasdataframegroup-by

(pandas) Create new column based on the presence or not of an element in groupby object


I have a problem that is close to this one :

(pandas) Create new column based on first element in groupby object

So I have this dataframe:

     accidentID      cartype   
0    58              70     
1    58              -70      
2    58              70     
3    58              100       
4    71              100   
5    71              -70    
6    250             70   
7    250             70  
8    250             100  
9    250             70  
10   70              70

What I want is to add a new column to show the car_in_the_wrong for a given accidentID knowing that :

if a cartype = -70 is present --> car_in_the_wrong = -70

if a cartype = 100 is present and no cartype = -70 --> car_in_the_wrong = 100

if 100 and -70 absent --> car_in_the_wrong = 70

     accidentID      cartype        car_in_the_wrong
0    58              70             -70
1    58              -70            -70
2    58              70             -70
3    58              100            -70
4    71              100            -70
5    71              -70            -70
6    250             70             100
7    250             70             100
8    250             100            100
9    250             70             100
10   300             70             70

Is there a way to do it with a groupby?

Thanks a lot!


Solution

  • Following this answer on how to Check if a value exists using multiple conditions within group in pandas, and taking advantage that statements with or are evaluated from left to right (i.e. the right-side part is not evaluated if the left-side holds True):

    df['car_in_the_wrong'] = df['accidentID'].map(df.groupby('accidentID').apply(lambda x: x['cartype'].eq(-70).any()*-70 or x['cartype'].eq(100).any()*100 or x['cartype'].eq(70).any()*70))