如何按配置指定前缀对DataFrame列分组并求和？

问题分析 #

你原来的代码里，列表推导式用了双重循环for col in cols for cat in granularity_suffix_list，这会让每个列名对应每个前缀都生成一个元素，导致分组键列表长度是原列数×前缀数，完全不符合分组要求。正确的逻辑应该是每个列只对应一个分组键：要么匹配到前缀就用前缀作为键，要么保留原列名。

正确实现方案 #

我们可以写一个函数，为每个列生成对应的分组键：遍历每个列名，检查是否以granularity_suffix_list中的任意前缀开头，找到第一个匹配的前缀作为分组键；如果没有匹配的前缀，就保留原列名。之后用这个分组键列表对DataFrame的列进行分组求和。

完整代码

import pandas as pd

# 配置参数
config = {
    'exclude_granularity': True,
    'granularity_suffix_list': ['A', 'B']
}

# 目标DataFrame
tt = pd.DataFrame({
    'A_2': [1,2,3],
    'A_3': [3,4,2],
    'B_4': [5,2,1],
    'B_1': [8,2,1],
    'C_3': [2,4,2]
})

if config['exclude_granularity']:
    def get_group_key(col, prefixes):
        # 遍历前缀列表，找到第一个匹配的前缀
        for prefix in prefixes:
            if col.startswith(prefix):
                return prefix
        # 没有匹配的前缀，返回原列名
        return col
    
    # 为每个列生成分组键
    group_keys = [get_group_key(col, config['granularity_suffix_list']) for col in tt.columns]
    # 按分组键求和
    result_df = tt.groupby(group_keys, axis=1).sum()

print(result_df)

输出结果

A   B  C_3
0  4  13    2
1  6   4    4
2  5   2    2

补充说明 #

• 如果需要支持包含指定字符串而不仅仅是开头匹配，只需要把col.startswith(prefix)改成prefix in col即可。
• 这个实现保证每个列只对应一个分组键，避免了原代码中重复生成键的问题，分组逻辑清晰且符合需求。

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