python - pandas groupby 与 if 条件

Question

我有以下带有发票数据的数据框。我想根据以下逻辑添加 MainCode。

1 - 第一个分组为 ticket_id 和 id。如果 level 为零，则 MainCode 应为零。如果不是 level 的 code，则应采用零。

+-----------+----+-------+------+
| 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 |
+-----------+----+-------+------+

到目前为止，我已经编写了以下代码

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

但我无法获得正确的输出。

我想要的输出如下

+-----------+----+-------+------+----------+
| 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 |
+-----------+----+-------+------+----------+

请指导我解决这个问题

Answer 1

您可以检查level中的哪些值与0不同，并将 bool 结果与相应组的第一个值相乘，可以通过groupby.transform和使用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