Python Pandas：按p_id分组计算指定b_id的val和并生成新列

现有如下DataFrame：

p_id    b_id      val
101     201       99
101     202       2 
101     203       11
102     201       1
102     202       50
102     203       23

需要按p_id分组，计算两个聚合值：

• sum_1：b_id为201和202的val之和
• sum_2：b_id为201和203的val之和

最终目标结果如下：

p_id   sum_1     sum_2
101    101        110
102     51         24

尝试了以下代码，但无法得到目标格式：

df2 = df.groupby(['p_id', 'b_id'])['val'].sum().reset_index()

请问该如何实现？

你当前的代码是按p_id和b_id双重分组，会保留每个b_id的单独求和结果，不符合需求。下面提供几种可行的实现方式：

方法一：分组后直接条件求和 #

通过groupby聚合时，使用lambda函数筛选符合条件的b_id对应的val求和：

result = df.groupby('p_id').agg(
    sum_1=('val', lambda x: x[df.loc[x.index, 'b_id'].isin([201, 202])].sum()),
    sum_2=('val', lambda x: x[df.loc[x.index, 'b_id'].isin([201, 203])].sum())
).reset_index()

方法二：先标记分组范围再聚合 #

先给每行添加布尔标识列，标记该行是否属于sum_1或sum_2的计算范围，再分组求和：

# 添加标识列
df['in_sum1'] = df['b_id'].isin([201, 202])
df['in_sum2'] = df['b_id'].isin([201, 203])

# 分组计算
result = df.groupby('p_id').agg(
    sum_1=('val', lambda x: (x * df.loc[x.index, 'in_sum1']).sum()),
    sum_2=('val', lambda x: (x * df.loc[x.index, 'in_sum2']).sum())
).reset_index()

# 可选：删除临时添加的标识列
df.drop(['in_sum1', 'in_sum2'], axis=1, inplace=True)

方法三：透视后列间求和 #

先将数据透视成宽表，再直接对指定列做加法运算：

# 透视成宽表，缺失值填充0
pivot_df = df.pivot(index='p_id', columns='b_id', values='val').fillna(0)

# 计算sum_1和sum_2
result = pivot_df.assign(
    sum_1=pivot_df[201] + pivot_df[202],
    sum_2=pivot_df[201] + pivot_df[203]
)[['sum_1', 'sum_2']].reset_index()

以上三种方法都能得到你想要的结果，可根据个人习惯选择。

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