Pandas多Label场景下分组聚合分析的优化实现方案问询

问题描述 #

现有演示数据集：

import pandas as pd
input=pd.DataFrame({'ID':['A','A','A','A','B','B','B','B','C','C','C','C'],
              'Year':[2000,2000,2000,2000,2000,2000,2000,2000,2001,2001,2001,2001],
               'Item':['a1','a2','a3','a4','b1','b2','b3','b4','c1','c2','c3','c4'],
                'Price':[1,3,4,5,2,4,7,3,5,7,6,1],
                'Label 1':[1,1,0,0,1,1,1,0,0,0,0,1],
                'Label 2':[0,1,1,0,0,1,1,0,1,0,0,1]})

需按ID和Year分组，输出如下聚合结果：

output=pd.DataFrame({'ID':['A','B','C'],
              'Year':[2000,2000,2001],
               'Label 1 Count':[2,3,1],
             'Label 2 Count':[2,2,2],
             'Label 1 Total Price':[4,13,1],
             'Label 2 Total Price':[7,11,6],
              'Item Total':[4,4,4]})

核心需求：

• 统计每个Label值为1的数量
• 计算对应Label为1时的Price总和
• 统计每组的Item总数

当前实现需手动创建Label X Price中间列，Label数量较多时效率极低，需批量处理方案。

解决方案 #

通过动态生成聚合规则，结合筛选定位Label列，无需手动创建中间列：

方案1：分步实现 #

1. 定位所有Label列

label_cols = input.filter(like='Label').columns.tolist()

2. 构建聚合字典
   为每个Label生成两组规则，同时加入Item计数：

agg_dict = {}
for col in label_cols:
    # 统计Label为1的数量
    agg_dict[col] = 'sum'
    # 计算Label为1时的总价
    agg_dict[f'{col} Total Price'] = lambda x, col=col: (x * input.loc[x.index, 'Price']).sum()

# 加入Item总数统计
agg_dict['Item'] = 'count'

3. 执行聚合并重命名列

result = input.groupby(['ID', 'Year']).agg(agg_dict).rename(columns={
    col: f'{col} Count' for col in label_cols
}).reset_index()

# 调整列顺序（可选，匹配期望输出）
desired_cols = ['ID', 'Year'] + [f'{col} Count' for col in label_cols] + [f'{col} Total Price' for col in label_cols] + ['Item']
result = result[desired_cols].rename(columns={'Item': 'Item Total'})

方案2：链式紧凑写法 #

利用字典解包直接生成聚合规则，代码更简洁：

label_cols = input.filter(like='Label').columns

result = (input.groupby(['ID', 'Year'])
          .agg(**{f'{col} Count': (col, 'sum') for col in label_cols},
               **{f'{col} Total Price': (col, lambda x: (x * input.loc[x.index, 'Price']).sum()) for col in label_cols},
               **{'Item Total': ('Item', 'count')})
          .reset_index())

两种方案均可自适应任意数量的Label列，输出结果与期望完全一致。

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