如何在Pandas DataFrame中插入分组求和行至每组开头？

问题描述 #

需要按name字段对DataFrame分组求和，并将求和行插入到对应分组的开头，示例如下：

原DataFrame：

name   place   eng_marks    math_marks
----------------------------------------------
0  prince  uk        70             80
1  prince  us        60             70
3  prince  fr        50             60
4  john    it        40             50
5  john    uk        30             40
6  lucy    np        20             10

期望结果：

name   place   eng_marks    math_marks
----------------------------------------------
0  prince  uk        180            210  ----> sum of groupby prince
1  prince  uk        70             80
2  prince  us        60             70
3  prince  fr        50             60
4  john    it        70             90   ----> sum of groupby john
5  john    it        40             50
6  john    uk        30             40
7  lucy    np        20             10   ----> sum of groupby lucy
8  lucy    np        20             10

高效解决方案 #

推荐使用矢量化操作实现，避免低效的循环或多次合并，以下是两种最优方案：

方案1：分组内拼接求和行（直观易读）

通过groupby.apply在每个分组内部先插入求和行，再保留原数据：

import pandas as pd

# 构造原数据
data = {
    'name': ['prince', 'prince', 'prince', 'john', 'john', 'lucy'],
    'place': ['uk', 'us', 'fr', 'it', 'uk', 'np'],
    'eng_marks': [70, 60, 50, 40, 30, 20],
    'math_marks': [80, 70, 60, 50, 40, 10]
}
df = pd.DataFrame(data)

def add_sum_row(group):
    # 复制组内第一行作为求和行模板，替换数值列为分组求和结果
    sum_row = group.iloc[0].copy()
    sum_row[['eng_marks', 'math_marks']] = group[['eng_marks', 'math_marks']].sum()
    # 拼接求和行与原分组数据
    return pd.concat([pd.DataFrame([sum_row]), group], ignore_index=True)

# 分组处理并合并结果
result_df = df.groupby('name', group_keys=False).apply(add_sum_row).reset_index(drop=True)

方案2：矢量化合并+排序（性能最优）

通过生成求和DataFrame，添加排序标识后与原数据合并排序，全程用Pandas矢量化操作，大数据量下效率更高：

import pandas as pd

# 构造原数据
data = {
    'name': ['prince', 'prince', 'prince', 'john', 'john', 'lucy'],
    'place': ['uk', 'us', 'fr', 'it', 'uk', 'np'],
    'eng_marks': [70, 60, 50, 40, 30, 20],
    'math_marks': [80, 70, 60, 50, 40, 10]
}
df = pd.DataFrame(data)

# 生成分组求和DataFrame，保留name和place的第一个值（与示例一致）
sum_df = df.groupby('name', as_index=False).agg(
    name='first',
    place='first',
    eng_marks=('eng_marks', 'sum'),
    math_marks=('math_marks', 'sum')
)

# 添加排序键：求和行排前面（0），原数据行排后面（1）
sum_df['sort_key'] = 0
df['sort_key'] = 1

# 合并后按name和sort_key排序，清理临时列
result_df = (pd.concat([sum_df, df])
             .sort_values(by=['name', 'sort_key'])
             .drop('sort_key', axis=1)
             .reset_index(drop=True))

方案优势 #

• 方案1逻辑直观，适合需要自定义分组内操作的场景；
• 方案2采用矢量化操作，避免了apply的隐式循环，在数据量较大时性能更优，是首选方案；
• 两种方案均无需多次合并操作，相比单独生成求和表再逐组合并的方式，减少了索引对齐和数据拷贝的开销。

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