Home > Software engineering >  Fill column based on conditional max value in Pandas
Fill column based on conditional max value in Pandas

Time:11-20

I have a dataframe that looks like this (link to csv):

id, time, value, approved
0,  0:00, 10,    false
1,  0:01, 20,    false
1,  0:02, 50,    false
1,  0:03, 20,    true
1,  0:04, 40,    true
1,  0:05, 40,    true
1,  0:06, 20,    false
2,  0:07, 35,    false
2,  0:08, 35,    false
2,  0:09, 50,    true
2,  0:10, 50,    true

and I want to compute a column that should be true for the first max approved value per ID. So it should be like this:

id, time, value, approved, is_max
0,  0:00, 10,    false,    false
1,  0:01, 20,    false,    false
1,  0:02, 50,    false,    false
1,  0:03, 20,    true,     false
1,  0:04, 40,    true,     true
1,  0:05, 40,    true,     false
1,  0:06, 20,    false,    false
2,  0:07, 35,    false,    false
2,  0:08, 35,    false,    false
2,  0:09, 50,    true,     true
2,  0:10, 50,    true,     false

I can achieve something close to this with

df['is_max'] = df['value'] == df.groupby(['id', df['approved']])['value'].transform('max').where(df['approved'])

but this will set to true both lines with a max value per ID (0:04 and 0:05 for ID 1 | 0:09 and 0:10 for ID 2). I just want the first row with the max value to be set to true.

CodePudding user response:

Here is an approach using pandas.DataFrame.mask based on your solution :

approved_1st_max = df.mask(~df["approved"]).groupby("id")["value"].transform('idxmax')

df["is_max"]= df.reset_index()["index"].eq(approved_1st_max)

# Output :

print(df)

    id  time  value  approved  is_max
0    0  0:00     10     False   False
1    1  0:01     20     False   False
2    1  0:02     50     False   False
3    1  0:03     20      True   False
4    1  0:04     40      True    True
5    1  0:05     40      True   False
6    1  0:06     20     False   False
7    2  0:07     35     False   False
8    2  0:08     35     False   False
9    2  0:09     50      True    True
10   2  0:10     50      True   False
  • Related