如何用Pandas groupby.filter()实现组内条件过滤移除指定行？

原始数据 #

import pandas as pd

df = pd.DataFrame({
    "Part Number" : [123456, 123456, 123456, 123456, 987654, 987654, 987654, 987654, 654321, 654321, 654321, 654321],
    "Operation" : [0, 100, 200, 300, 0, 100, 200, 300, 0, 100, 200, 300],
    "Source" : ["A", "A", "A", "A", "A", "B", "A", "A", "B", "B", "B", "B"]
})

需求说明 #

按Part Number分组后：

• 若组内存在Source=="A"的行，移除该组所有Source=="B"的行
• 若组内无"A"，保留全部行

此前尝试用df.groupby("Part Number").filter(lambda x: (x["Source"] == "A").count() > 0)未达预期，且不想用循环实现。

预期输出 #

pd.DataFrame({
    "Part Number" : [123456, 123456, 123456, 123456, 987654, 987654, 987654, 654321, 654321, 654321, 654321],
    "Operation" : [0, 100, 200, 300, 0, 200, 300, 0, 100, 200, 300],
    "Source" : ["A", "A", "A", "A", "A", "A", "A", "B", "B", "B", "B"]
})

解决方案 #

groupby().filter()的作用是保留或删除整个分组，无法实现仅删除分组内部分行的需求。可以通过groupby.transform()生成标记列，再结合布尔索引实现需求：

# 标记每个行所在分组是否存在Source="A"
has_a = df.groupby("Part Number")["Source"].transform(lambda x: (x == "A").any())

# 构造过滤条件：分组无A则全部保留；分组有A则仅保留Source为A的行
filtered_df = df[(~has_a) | (df["Source"] == "A")]

执行后filtered_df与预期输出完全一致。

内容的提问来源于stack exchange，提问作者smichael_44