Tommy
Tommy

Reputation: 515

pandas groupby with if condition

I have the below dataframe with invoice data. I want to add the MainCode, based on the below logic.

1 - first groupby ticket_id & id. If the level is zero, then the MainCode should be zero. if not code of the level zero should taken.

+-----------+----+-------+------+
| ticket_id | id | level | code |
+-----------+----+-------+------+
|         1 |  0 |     0 | 1710 |
|         1 |  0 |     1 |  372 |
|         1 |  0 |     2 |  607 |
|         1 |  1 |     0 | 1727 |
|         1 |  1 |     1 |  370 |
|         1 |  1 |     2 |  607 |
|         2 |  0 |     0 |  269 |
|         2 |  0 |     1 |  371 |
|         2 |  0 |     2 |  607 |
|         2 |  1 |     0 |  277 |
|         2 |  1 |     1 |  371 |
|         2 |  1 |     2 |  607 |
+-----------+----+-------+------+

So far, I have written the below code

df.groupby(['ticket_id','id'])['code'].transform(lambda x: if df['level'] == 0, 0, df['code']).

but I'm not able to get the correct out put.

my desired output is as below

+-----------+----+-------+------+----------+
| ticket_id | id | level | code | MainCode |
+-----------+----+-------+------+----------+
|         1 |  0 |     0 | 1710 |        0 |
|         1 |  0 |     1 |  372 |     1710 |
|         1 |  0 |     2 |  607 |     1710 |
|         1 |  1 |     0 | 1727 |        0 |
|         1 |  1 |     1 |  370 |     1727 |
|         1 |  1 |     2 |  607 |     1727 |
|         2 |  0 |     0 |  269 |        0 |
|         2 |  0 |     1 |  371 |      269 |
|         2 |  0 |     2 |  607 |      269 |
|         2 |  1 |     0 |  277 |        0 |
|         2 |  1 |     1 |  371 |      277 |
|         2 |  1 |     2 |  607 |      277 |
+-----------+----+-------+------+----------+

please guide me to solve this

Upvotes: 1

Views: 237

Answers (1)

yatu
yatu

Reputation: 88305

You could check which values in level are different to 0, and multiply the boolean result with the first value of the corresponding group, which can obtained taking the groupby.transform and aggregating with first:

df['MainCode'] = (df.level.ne(0)
                    .mul(df.groupby(['ticket_id','id']).code
                    .transform('first')))

    ticket_id  id  level  code  MainCode
0           1   0      0  1710         0
1           1   0      1   372      1710
2           1   0      2   607      1710
3           1   1      0  1727         0
4           1   1      1   370      1727
5           1   1      2   607      1727
6           2   0      0   269         0
7           2   0      1   371       269
8           2   0      2   607       269
9           2   1      0   277         0
10          2   1      1   371       277
11          2   1      2   607       277

Upvotes: 2

Related Questions