基于列头共性合并Pandas DataFrame列的技术求助

我有一个包含事件和日期的DataFrame，希望解析列头，将具有共性的列合并为一列，得到指定输出。目前尝试的代码未达到预期效果（前两列存在缺失值），仅'dis'列结果正确，寻求解决方法。

输入输出示例 #

输入DataFrame #

import pandas as pd

df1 = pd.DataFrame({'yyyyww': ['2022-01','2022-02','2022-03', '2022-01','2022-02','2022-03','2022-01','2022-03'],
                         '2001_52': [10,0,0,0,0,0,0,0],
                         '2002_52': [0,30,0,0,0,0,0,0],
                         '2003_52': [0,0,50,0,0,0,0,0],
                         '2001_23': [0,0,0,20,0,0,0,0],
                         '2002_23': [0,0,0,0,15,0,0,0],
                         'dis20': [0,0,0,0,0,25,0,0],
                         'dis30': [0,0,0,0,0,0,15,0],
                         'dis40': [0,0,0,0,0,0,0,75],})

预期输出DataFrame #

df2 = pd.DataFrame({'yyyyww': ['2022-01','2022-02','2022-03', '2022-01','2022-02','2022-03','2022-01','2022-03'],
                         'wk_52': [10,30,50,0,0,0,0,0],
                         'wk_23': [0,0,0,20,15,0,0,0],
                         'dis': [0,0,0,0,0,25,15,75],})

现有尝试及问题 #

尝试代码 #

# Create empty dict 
merge_data = {}

# loop through columns to parse name 
for col in df1.columns:
    if col.startswith('2') and '_' in col:
        key = col.split('_')[-1]
        if key.isdigit():
            if key not in merge_data:
                merge_data[key] = df1[col]
        else:
            merge_data[key] +=df1[col]

# merge columns and apply values to empty dict
merge_data_out = {f"wk_{key}": values for key, values in merge_data.items()}

# procedure for combining 'dis...' columns
merge_data_out['dis'] = df1.filter(like='dis').sum(axis=1)
merge_data_out = pd.DataFrame(merge_data_out)
merge_data_out

问题说明 #

运行后wk_52和wk_23列仅保留了第一个匹配列的数据，其他行未完成累加，不符合预期，仅dis列结果正确。

问题分析与解决方案 #

问题根源 #

循环逻辑存在漏洞：当同key的后续列出现时，代码未执行累加操作——else分支仅针对key非数字的情况，而此处key均为数字，导致同组列无法合并。

修正方案 #

方法一：修复循环逻辑

import pandas as pd

df1 = pd.DataFrame({'yyyyww': ['2022-01','2022-02','2022-03', '2022-01','2022-02','2022-03','2022-01','2022-03'],
                         '2001_52': [10,0,0,0,0,0,0,0],
                         '2002_52': [0,30,0,0,0,0,0,0],
                         '2003_52': [0,0,50,0,0,0,0,0],
                         '2001_23': [0,0,0,20,0,0,0,0],
                         '2002_23': [0,0,0,0,15,0,0,0],
                         'dis20': [0,0,0,0,0,25,0,0],
                         'dis30': [0,0,0,0,0,0,15,0],
                         'dis40': [0,0,0,0,0,0,0,75],})

merge_data = {}

# 修正循环：所有同key列执行累加
for col in df1.columns:
    if col.startswith('2') and '_' in col:
        key = col.split('_')[-1]
        if key.isdigit():
            if key not in merge_data:
                # 初始化全0Series避免缺失
                merge_data[key] = pd.Series([0]*len(df1), index=df1.index)
            merge_data[key] += df1[col]

merge_data_out = {f"wk_{key}": values for key, values in merge_data.items()}
merge_data_out['dis'] = df1.filter(like='dis').sum(axis=1)
merge_data_out['yyyyww'] = df1['yyyyww']

# 调整列顺序匹配预期
df2 = pd.DataFrame(merge_data_out)[['yyyyww', 'wk_52', 'wk_23', 'dis']]
print(df2)

方法二：Pandas原生方法（推荐）

利用melt+groupby实现，代码更简洁易维护：

import pandas as pd

df1 = pd.DataFrame({'yyyyww': ['2022-01','2022-02','2022-03', '2022-01','2022-02','2022-03','2022-01','2022-03'],
                         '2001_52': [10,0,0,0,0,0,0,0],
                         '2002_52': [0,30,0,0,0,0,0,0],
                         '2003_52': [0,0,50,0,0,0,0,0],
                         '2001_23': [0,0,0,20,0,0,0,0],
                         '2002_23': [0,0,0,0,15,0,0,0],
                         'dis20': [0,0,0,0,0,25,0,0],
                         'dis30': [0,0,0,0,0,0,15,0],
                         'dis40': [0,0,0,0,0,0,0,75],})

# 处理wk_*列
wk_cols = [col for col in df1.columns if col.startswith('2') and '_' in col]
wk_df = df1.melt(id_vars='yyyyww', value_vars=wk_cols, var_name='col', value_name='value')
wk_df['wk'] = wk_df['col'].str.split('_').str[-1]
wk_sum = wk_df.groupby(['yyyyww', 'wk'])['value'].sum().unstack().add_prefix('wk_')

# 处理dis列
dis_sum = df1.filter(like='dis').sum(axis=1).rename('dis')

# 合并结果并补0
df2 = pd.concat([df1['yyyyww'], wk_sum, dis_sum], axis=1).fillna(0)
# 调整列顺序
df2 = df2[['yyyyww', 'wk_52', 'wk_23', 'dis']]
print(df2)

输出验证 #

两种方法均会输出与预期一致的结果：

yyyyww  wk_52  wk_23  dis
0  2022-01   10.0    0.0    0
1  2022-02   30.0    0.0    0
2  2022-03   50.0    0.0    0
3  2022-01    0.0   20.0    0
4  2022-02    0.0   15.0    0
5  2022-03    0.0    0.0   25
6  2022-01    0.0    0.0   15
7  2022-03    0.0    0.0   75

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