Pandas groupby drop_duplicates 基于多列上的多个条件

Question

我有一个这样的数据集：

ID    Data    AddType    Num
123   What    HA1        1
123   I       HA1        2
123   What    OT1        1
456   I       OT2        1
456   Gotta   OT2        2

这样，我需要按 ID 分组，并根据 AddType 和 Num 的两条规则保留数据：

1. 如果存在一个 HA1 类型行，请保留它
2. 如果不存在 HA1 类型行，则保留 Num 值最高的行
3. 如果存在多个 HA1 类型行，请保留 Num 值最高的行

所以基本上优先级 1 是有一个 HA1 类型的行。第二个优先级（如果有多个 HA1 类型行或没有），优先级 2 是 Num 值。所以期望的结果是：

ID    Data    AddType    Num
123   I       HA1        2
456   Gotta   OT2        2

我已经尝试了 gropuby 和 drop_duplicates 的几种变体，但我似乎无法弄清楚这一点。我怎样才能做到这一点？谢谢！

Answer 1

试试sort_values 然后drop_duplicates

out = df.sort_values(by = 'AddType', key = lambda x : x.eq('HA1')).sort_values('Num').drop_duplicates('ID',keep='last')
Out[506]: 
    ID   Data AddType  Num
4  456  Gotta     OT2    2
1  123      I     HA1    2