如何用Pandas的groupby或pivot_table实现基于列过滤的数据聚合？

现有DataFrame #

col1  col2  col3  col4  col5
A     A1    X     1      2
A     A2    Y     2      2
A     A3    Z     1      2
B     B1    X     2      2
B     B2    Y     2      2
B     B3    Z     1      2

映射规则 #

mapping：

col1  col3
 A --> X
 A --> Y
(A,B) --> Z
(A,B) --> (Y,Z)

规则说明：

• A-->X：筛选col1为A且col3为X的数据，计算col4的求和值
• (A,B)-->Z：筛选col1为A或B且col3为Z的数据，计算col4的求和值

尝试代码 #

df1 = df.pivot_table(index= [col1,col3], columns = col3, values = col4, aggfunc='sum')

期望结果 #

col1  X   Y   Z   YZ
A     1   2   2   6
B             2   6

请问是否可以使用groupby或pivot_table实现该需求？

单纯用groupby或pivot_table无法直接覆盖所有规则场景，需要先结合映射规则生成自定义分组标识，再配合聚合操作实现目标。

具体步骤： #

1. 结构化映射规则：将规则转为便于处理的字典格式

rules = [
    {"col1": {"A"}, "col3": {"X"}, "name": "X"},
    {"col1": {"A"}, "col3": {"Y"}, "name": "Y"},
    {"col1": {"A", "B"}, "col3": {"Z"}, "name": "Z"},
    {"col1": {"A", "B"}, "col3": {"Y", "Z"}, "name": "YZ"},
]

2. 按规则筛选并聚合：遍历每条规则，筛选符合条件的数据，按col1分组求和

import pandas as pd

# 假设原始数据已存入DataFrame df
result_list = []
for rule in rules:
    # 生成筛选掩码
    filter_mask = df["col1"].isin(rule["col1"]) & df["col3"].isin(rule["col3"])
    # 筛选后按col1分组求和
    agg_result = df[filter_mask].groupby("col1")["col4"].sum().reset_index()
    agg_result["metric"] = rule["name"]
    result_list.append(agg_result)

# 合并所有规则的聚合结果
combined_df = pd.concat(result_list)

3. 转成宽表匹配期望格式：将长表转成宽表，填充空值并调整列顺序

final_df = combined_df.pivot(index="col1", columns="metric", values="col4").fillna("").reset_index()
final_df.columns.name = None
# 调整列顺序与期望结果一致
final_df = final_df[["col1", "X", "Y", "Z", "YZ"]]

执行后final_df的结果就与你期望的完全一致。这种方法逻辑清晰，每条规则对应明确的处理步骤，后续调整规则也很方便。

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