如何用groupby与聚合函数重塑Pandas DataFrame？

问题需求 #

现有如下结构的Pandas DataFrame：

A      B     C     D            E      F
foo    s     HO    02/01/24     100    20.0
foo    s     HO    02/01/24     200    20.0
foo    b     HO    02/01/24     100    20.0
foo    b     HO    02/01/24     200    20.0

bar    s     HO    02/01/24     100    20.0
bar    s     HO    02/01/24     200    20.0
bar    b     HO    02/01/24     100    20.0
bar    b     HO    02/01/24     200    20.0

fizz    s     BRT    02/01/24     100    20.0
fizz    s     BRT    02/01/24     200    20.0
fizz    b     BRT    02/01/24     100    20.0
fizz    b     BRT    02/01/24     200    20.0

buzz    s     PW    02/01/24     100    20.0
buzz    s     PW    02/01/24     200    20.0
buzz    b     PW    02/01/24     100    20.0
buzz    b     PW    02/01/24     200    20.0

需要完成以下操作：

• 按A、B、C、D列分组
• 计算每组E列的总和，提取每组F列的第一个值
• 将结果重塑为指定的层级索引格式

输入数据代码 #

import pandas as pd

data = {
    'A': ['foo', 'foo', 'foo', 'foo', 'bar', 'bar', 'bar', 'bar', 'fizz', 'fizz', 'fizz', 'fizz', 'buzz', 'buzz', 'buzz', 'buzz'],
    'B': ['s', 's', 'b', 'b', 's', 's', 'b', 'b', 's', 's', 'b', 'b', 's', 's', 'b', 'b'],
    'C': ['HO', 'HO', 'HO', 'HO', 'HO', 'HO', 'HO', 'HO', 'BRT', 'BRT', 'BRT', 'BRT', 'PW', 'PW', 'PW', 'PW'],
    'D': ['02/01/24', '02/01/24', '02/01/24', '02/01/24', '02/01/24', '02/01/24', '02/01/24', '02/01/24', '02/01/24', '02/01/24', '02/01/24', '02/01/24', '02/01/24', '02/01/24', '02/01/24', '02/01/24'],
    'E': [100, 200, 100, 200, 100, 200, 100, 200, 100, 200, 100, 200, 100, 200, 100, 200],
    'F': [20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0]
}
df = pd.DataFrame(data)

用户尝试的代码 #

df_groups = df.groupby(['A','C','B','D'], sort=False, as_index=False).agg({'E': 'sum', 'F': 'first'})
df_list = [x for _, x in df_groups.groupby(['A', 'B'])]

解决步骤与代码 #

步骤1：分组聚合 #

调整分组顺序为['B', 'C', 'A', 'D']，完成E求和、F取首值的聚合：

# 分组聚合，保留索引以便后续处理
df_agg = df.groupby(['B', 'C', 'A', 'D'], sort=False).agg(
    E_sum=('E', 'sum'),
    F_first=('F', 'first')
).reset_index()

步骤2：调整结构与索引 #

设置多级索引并调整列顺序，匹配目标格式：

# 设置B、C为多级索引
df_result = df_agg.set_index(['B', 'C'])
# 调整列顺序并重命名
df_result = df_result[['D', 'A', 'E_sum', 'F_first']].rename(columns={'E_sum': 'E', 'F_first': 'F'})
# 可选：将F列转为整数，匹配目标格式的20而非20.0
df_result['F'] = df_result['F'].astype(int)

步骤3：优化显示效果 #

开启Pandas的多级索引稀疏显示，让重复索引值隐藏：

with pd.option_context('display.multi_sparse', True):
    print(df_result)

最终输出：

D       A    E   F
B   C                        
b   HO  02/01/24  foo  300  20
        02/01/24  bar  300  20
    BRT 02/01/24  fizz 300  20
    PW  02/01/24  buzz 300  20
s   HO  02/01/24  foo  300  20
        02/01/24  bar  300  20
    BRT 02/01/24  fizz 300  20
    PW  02/01/24  buzz 300  20

说明 #

• 调整分组顺序是为了后续设置多级索引更直接，也可以用原分组顺序后重新排列索引。
• display.multi_sparse选项用于简化多级索引的显示，让输出更贴近目标格式。

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