熊猫groupby与if条件

Question

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

Answer 1

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