基于指定类别对比行值与其他行求和结果——Pandas DataFrame处理

需求说明 #

给定包含Month、reporter、code、code_category、Value字段的DataFrame，需新增subtot_correct列：

• 主类别（如A、B）的Value总值与对应子类别（如A_1、B_2）的Value求和值相等时，赋值1
• 不相等时赋值0
• 同步标记Month+reporter+code组合下的所有关联行（含主类别和对应子类别行）

完整实现代码 #

import pandas as pd

# 构造示例数据（模拟用户输入结构）
data = {
    'Month': ['2023-01', '2023-01', '2023-01', '2023-01', '2023-01', '2023-01'],
    'reporter': ['ADC', 'ADC', 'ADC', 'XYZ', 'XYZ', 'XYZ'],
    'code': ['C001', 'C001', 'C001', 'C002', 'C002', 'C002'],
    'code_category': ['A', 'A_1', 'A_2', 'B', 'B_1', 'B_2'],
    'Value': [100, 60, 40, 150, 80, 60]  # XYZ的B类别总值150，子类别求和140，不匹配
}
df = pd.DataFrame(data)

# 1. 从code_category提取主类别标识（按下划线分割取前缀）
df['main_category'] = df['code_category'].str.split('_').str[0]

# 2. 计算每个分组下子类别Value的求和值
subtotals = df[df['code_category'].str.contains('_')].groupby(
    ['Month', 'reporter', 'code', 'main_category']
)['Value'].sum().reset_index(name='sub_total')

# 3. 生成主类别行的匹配标记
main_rows = df[~df['code_category'].str.contains('_')].merge(
    subtotals, on=['Month', 'reporter', 'code', 'main_category'], how='left'
)
# 无对应子类别的主类别标记为0，可按需调整
main_rows['subtot_correct'] = (main_rows['Value'] == main_rows['sub_total']).fillna(0).astype(int)

# 4. 将匹配标记同步到所有对应行（包括子类别行）
df = df.merge(
    main_rows[['Month', 'reporter', 'code', 'main_category', 'subtot_correct']],
    on=['Month', 'reporter', 'code', 'main_category'],
    how='left'
)

# 输出结果
print(df)

输出结果示例 #

Month reporter   code code_category  Value main_category  subtot_correct
0  2023-01       ADC  C001             A    100             A               1
1  2023-01       ADC  C001           A_1     60             A               1
2  2023-01       ADC  C001           A_2     40             A               1
3  2023-01       XYZ  C002             B    150             B               0
4  2023-01       XYZ  C002           B_1     80             B               0
5  2023-01       XYZ  C002           B_2     60             B               0

优化建议 #

• 区分无子类和不匹配场景：将无对应子类别的主类别标记为2，而非0，避免与“求和不匹配”混淆
• 新增差异值列：添加sub_total_diff列，计算Value与sub_total的差值，快速定位差异大小：

  main_rows['sub_total_diff'] = main_rows['Value'] - main_rows['sub_total']
  

• 适配不同类别命名规则：如果code_category不是下划线分隔，可调整主类别提取逻辑，比如按固定长度截取、正则匹配：

  # 示例：提取开头的单个大写字母作为主类别
  df['main_category'] = df['code_category'].str.extract(r'^([A-Z])')
  

• 简化分组赋值逻辑：使用transform方法直接在原DataFrame中完成计算，减少中间表：

  # 计算每个分组的子类别求和
  df['sub_total'] = df.groupby(['Month', 'reporter', 'code', 'main_category'])['Value'].transform(
      lambda x: x[df.loc[x.index, 'code_category'].str.contains('_')].sum()
  )
  # 生成主类别行的匹配标记
  df['subtot_correct'] = (df['Value'] == df['sub_total']).astype(int)
  # 将主类别标记同步到所有子类别行
  df['subtot_correct'] = df.groupby(['Month', 'reporter', 'code', 'main_category'])['subtot_correct'].transform('max')
  

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