动态列名DataFrame按年月分组求和并新增对应汇总列

核心思路 #

先按列名的「年月」前缀分组，计算每组三列的横向求和，再将求和列插入到对应年月三列的后方。全程通过前缀匹配和索引定位实现，无需硬编码列名。

实现步骤及代码 #

1. 按年月前缀分组列名
   利用列名的_分隔符提取年月前缀，将同前缀的三列归为一组：

   from itertools import groupby
   
   # 提取列名中的年月前缀
   def extract_month(col_name):
       return col_name.split('_')[0]
   
   # 按年月前缀分组列（依赖原始列顺序为连续的同年月列）
   grouped_columns = {
       month: list(cols) 
       for month, cols in groupby(df1.columns, key=extract_month)
   }
   

2. 倒序处理分组并插入求和列
   倒序遍历分组避免插入列导致的索引偏移问题，计算每组求和后插入到对应位置：

   # 倒序遍历年月分组，确保插入位置不受前面操作影响
   for month in reversed(grouped_columns.keys()):
       target_cols = grouped_columns[month]
       # 计算当前年月三列的行求和
       sum_result = df1[target_cols].sum(axis=1, skipna=True)
       # 获取该组最后一列的索引，求和列插入到其下一位
       insert_index = df1.columns.get_loc(target_cols[-1]) + 1
       # 插入求和列
       df1.insert(insert_index, f"{month}_sum", sum_result)
   

关键注意事项 #

• 确保CSV读取后的列顺序是连续的同年月列（即「年月_a」「年月_b」「年月_c」依次排列），groupby才能正确分组；若列顺序混乱，需先按前缀排序：

  # 若列顺序混乱，先按年月前缀排序
  df1 = df1[sorted(df1.columns, key=extract_month)]
  

• 使用skipna=True（默认值）可自动跳过空值求和，若需严格不跳过空值可改为skipna=False。

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