dimanche 23 août 2020

If-else statement with group_by in Pandas dataframe

I’ve a pd df consists four columns: ID, t, x1 and x2.

import pandas as pd
dat = {'ID': [1,1,1,1,2,2,2,3,3,3,3,4,4,4,5,5,6,6,6],
        't': [0,1,2,3,0,1,2,0,1,2,3,0,1,2,0,1,0,1,2],
        'x1' : [3.5,3.5,3.5,3.5,2.01,2.01,2.01,3.9,3.9,3.9,3.9,2.2,2.2,2.2,1.8,1.8,2.1,2.1,2.1],
       'x2': [4,4,4,4,3,3,3,4,4,4,4,3,3,3,2,2,3,3,3]
        }

df = pd.DataFrame(dat, columns = ['ID', 't', 'x1','x2'])

print (df)

I need to create a new column y and group_by ID such that

if t!=max(t) then y=1,
if t==max(t) then y = x1-x2+1.

The output would look like:

enter image description here

Please not that I have million of records, so the faster the solution the better.

Aucun commentaire:

Enregistrer un commentaire