如何对转置数据集按条件聚合动态生成NRev系列新列

核心思路 #

把宽格式的配对Return/Rev列转为长格式处理，自动适配列数和Return取值范围，无需硬编码任何列名或判断逻辑。

完整代码 #

import pandas as pd
import numpy as np

# 假设df是你读入的原始DataFrame
# 第一步：自动识别所有Return和Rev列，保证顺序一一对应
return_cols = sorted([col for col in df.columns if col.startswith('Return')], 
                     key=lambda x: int(x.replace('Return', '')))
rev_cols = sorted([col for col in df.columns if col.startswith('Rev')], 
                  key=lambda x: int(x.replace('Rev', '')))

# 第二步：转为长表结构，跳过空值
long_records = []
for _, row in df.iterrows():
    current_date = row['Date']
    for r_col, rev_col in zip(return_cols, rev_cols):
        r_val = row[r_col]
        rev_val = row[rev_col]
        if pd.isna(r_val) or pd.isna(rev_val):
            continue
        long_records.append({
            'Date': current_date,
            'return_val': int(r_val),
            'revenue': rev_val
        })
df_long = pd.DataFrame(long_records)

# 第三步：按Date和Return取值分组聚合营收
df_agg = df_long.groupby(['Date', 'return_val'], as_index=False)['revenue'].sum()

# 第四步：转回宽表，生成NRev系列列
df_pivot = df_agg.pivot(index='Date', columns='return_val', values='revenue').fillna(0).astype(int)
# 补全从0到最大Return值之间的所有列，缺失值填0
max_return = int(df[return_cols].max().max())
full_cols = [f'NRev{i}' for i in range(max_return + 1)]
df_pivot.columns = [f'NRev{col}' for col in df_pivot.columns]
result = df_pivot.reindex(columns=full_cols, fill_value=0).reset_index()

方案优势 #

• 无需硬编码列名：新增Return/Rev配对列会自动识别，无需修改代码
• 自动适配取值范围：Return最大值变化时自动生成对应数量的NRev列，无需手动添加判断逻辑
• 原生支持Date重复值聚合：同一Date下的同Return取值营收会自动累加，符合分组要求
• 自动处理空值：Return或Rev为nan的行自动跳过，无需额外清洗操作

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