如何基于分组多列条件筛选Pandas DataFrame？（附业务场景）

问题场景 #

给定如下Pandas DataFrame：

import pandas as pd

data = [
    ["AUD", "A", "STOCK", "$10.00"],
    ["AUD", "A", "BOND", "$10.00"],
    ["AUD", "B", "OPTION", "$11.00"],
    ["AUD", "B", "STOCK", "$12.00"],
    ["USD", "A", "STOCK", "$14.00"],
    ["USD", "A", "BOND", "$11.00"],
    ["USD", "A", "OPTION", "$19.00"],
    ["USD", "B", "BOND", "$12.00"]
]

df = pd.DataFrame(data, columns=["currency", "index", "product", "price"])

需求规则 #

按currency与index的组合分组：

• 若分组内包含OPTION类型产品，仅保留该分组中的OPTION行，过滤掉STOCK和BOND行
• 若分组内无OPTION类型产品，保留该分组所有行

预期输出 #

currency index product   price
0      AUD     A   STOCK  $10.00
1      AUD     A    BOND  $10.00
2      AUD     B  OPTION  $11.00
3      USD     A  OPTION  $19.00
4      USD     B    BOND  $12.00

解决方案 #

方法一：使用transform生成组标记筛选 #

通过groupby.transform将组级判断结果广播到每一行，再结合布尔索引筛选：

# 标记每行所在分组是否包含OPTION
has_option = df.groupby(["currency", "index"])["product"].transform(lambda x: x.eq("OPTION").any())

# 执行筛选：组内无OPTION则全留，组内有OPTION则仅留OPTION行
filtered_df = df[(~has_option) | (df["product"] == "OPTION")].reset_index(drop=True)

print(filtered_df)

方法二：自定义分组过滤函数 #

通过groupby.apply传入自定义函数，对每个分组单独处理：

def filter_group(group):
    # 判断分组内是否存在OPTION
    if "OPTION" in group["product"].values:
        return group[group["product"] == "OPTION"]
    else:
        return group

# 应用分组过滤并重置索引
filtered_df = df.groupby(["currency", "index"]).apply(filter_group).reset_index(drop=True)

print(filtered_df)

两种方法均能得到符合预期的筛选结果，可根据个人习惯选择。

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