如何根据pandas DataFrame同组汇总行的值为明细行新增总可用预算列

实现方法 #

你可以用pandas内置的where+分组transform功能实现，不需要遍历全表，代码如下：

import numpy as np
import pandas as pd

# 你的原始DataFrame构造代码
df = pd.DataFrame([
    [50000,'01',991,5000],
    [50000,'01',992,5000],
    [50000,'01','tot',10000],
    [50000,'02',777,2500],
    [50000,'02','tot',2500],
    [66666,'01',777,7000],
    [66666,'01',788,8000],
    [66666,'01','tot',15000]],
    columns=['ID','group_code','section_code','budget'])

# 核心逻辑
# 1. 提取每个(ID, group_code)分组的tot总预算，填充到同组所有行
df['total available'] = df['budget'].where(df['section_code'] == 'tot')\
    .groupby([df['ID'], df['group_code']]).transform('max')
# 2. 汇总行的total available设为NaN
df.loc[df['section_code'] == 'tot', 'total available'] = np.nan

逻辑说明 #

1. df['budget'].where(df['section_code'] == 'tot')会保留section_code='tot'行的预算值，其余行的取值设为NaN
2. 按ID和group_code分组后调用transform('max')，会将分组内唯一的非NaN值（即分组总预算）填充到分组内所有行
3. 最后单独将section_code='tot'的汇总行的新列值替换为NaN，即可得到你需要的结果

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