如何高效实现基于number与type的两类金额聚合列生成？

问题背景 #

原始数据集：

需要生成目标汇总表：

要求：

• amount A：按number分组后，仅统计type='A'的amount总和，无对应记录时填0
• amount A+B：按number分组后所有type的amount总和

当前仅能通过拆分多个子集分别计算两列，寻求更高效的实现方式。

高效实现方案（Python Pandas） #

无需拆分多个子集，可通过单次分组聚合完成两列的计算，以下提供两种简洁写法：

方法1：分组时直接嵌入条件求和 #

import pandas as pd

# 构造原始数据
df = pd.DataFrame({
    'number': [1,1,2,3,2,1],
    'type': ['A','A','A','B','B','B'],
    'amount': [10,20,10,20,10,20]
})

# 一次分组完成双列计算
result = df.groupby('number').agg(
    **{
        'amount A': ('amount', lambda x: x[df.loc[x.index, 'type'] == 'A'].sum()),
        'amount A+B': ('amount', 'sum')
    }
).reset_index()

# 无A类型的分组填充0
result['amount A'] = result['amount A'].fillna(0)

print(result)

方法2：先标记A类型金额再分组求和 #

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'number': [1,1,2,3,2,1],
    'type': ['A','A','A','B','B','B'],
    'amount': [10,20,10,20,10,20]
})

# 标记仅A类型的金额，非A设为0
df['A_amount'] = np.where(df['type'] == 'A', df['amount'], 0)

# 单次分组聚合两列
result = df.groupby('number').agg(
    **{
        'amount A': ('A_amount', 'sum'),
        'amount A+B': ('amount', 'sum')
    }
).reset_index()

print(result)

两种方法均通过单次分组操作完成计算，避免了多次子集拆分的冗余步骤，执行效率更高，且代码可读性更强。

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