Pandas多字段分组：按相似前缀分组并分类求和

需求说明 #

需要按name列的冒号前缀（如AA:3400提取AA）、type列的分类值（包含空值）以及month字段进行分组，对size列求和。

原始数据 #

name      type    size   month
AA:3400            5     august
AA:3401   FALSE    1     august
AA:3402   FALSE    2     august
AA:3404   TRUE     0     august
AA:3409   FALSE    1     september
AA:3410   FALSE    8     september
AA:3412   FALSE    9     september
BB:3400   TRUE     4     august
BB:3401   FALSE    7     august

期望结果 #

name    type    size  month
AA      TRUE    0     august
AA      FALSE   3     august
AA              5     august
BB      TRUE    4     august
BB      FALSE   7     august
AA      TRUE    0     september
AA      FALSE   18    september

当前代码的问题 #

1. 错误引用了不存在的date列，仅需处理name列即可
2. 分组时未加入month字段，导致结果丢失月份维度
3. type列的类型转换逻辑错误，不应基于name列转换，需保留原type的空值及分类值，并设置正确的分类顺序

正确实现代码 #

import pandas as pd

# 构造原始数据
data = {
    "name": ["AA:3400", "AA:3401", "AA:3402", "AA:3404", "AA:3409", "AA:3410", "AA:3412", "BB:3400", "BB:3401"],
    "type": ["", "FALSE", "FALSE", "TRUE", "FALSE", "FALSE", "FALSE", "TRUE", "FALSE"],
    "size": [5, 1, 2, 0, 1, 8, 9, 4, 7],
    "month": ["august", "august", "august", "august", "september", "september", "september", "august", "august"]
}
df = pd.DataFrame(data)

# 核心处理逻辑
out = (
    df
    # 提取name的前缀作为分组用的name字段
    .assign(name_prefix=df["name"].str.split(":").str[0])
    # 将type转换为包含空值的分类类型，指定顺序匹配期望结果
    .assign(type=pd.Categorical(df["type"], categories=["TRUE", "FALSE", ""], ordered=True))
    # 按前缀、type、month分组，dropna=False保留空值分组
    .groupby(["name_prefix", "type", "month"], dropna=False)
    ["size"].sum()
    .reset_index()
    # 恢复name列名并调整列顺序
    .rename(columns={"name_prefix": "name"})
    [["name", "type", "size", "month"]]
)

print(out)

代码说明 #

• name.str.split(":").str[0]：精准提取冒号前的前缀字符串，作为分组的name维度
• 分类类型设置：指定["TRUE", "FALSE", ""]的顺序，确保分组结果的排序和期望一致
• 分组字段包含month：保证数据按月份维度拆分统计
• dropna=False：保留type列为空值的分组，对应原始数据中AA:3400的单独统计行

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