按指定年份区间对DataFrame行数据求和生成新表的实现问题

以下代码默认满足「同一ID对应唯一code」的业务逻辑，和需求匹配。 #

实现步骤 #

• 导入依赖库

import pandas as pd

• 给原始DataFrame新增年份区间标签
  使用pd.cut方法按指定区间拆分年份：

df['year_period'] = pd.cut(
    df['Year'],
    bins=[2017, 2019, 2022], # 左开右闭区间，刚好覆盖2018-2019、2020-2022
    labels=['2018-2019', '2020-2022']
)

• 分组求和
  按ID、code、年份区间三个维度分组，对count_1、count_2字段求和：

grouped_df = df.groupby(
    ['ID', 'code', 'year_period'],
    as_index=False
)[['count_1', 'count_2']].sum()

• 转换为宽表格式
  将年份区间从行维度转为列维度，得到每个ID对应一行的最终结果：

final_df = grouped_df.pivot(
    index=['ID', 'code'],
    columns='year_period',
    values=['count_1', 'count_2']
).reset_index()
# 合并列名，方便后续使用
final_df.columns = [
    f'{col[0]}_{col[1]}' if col[1] else col[0] 
    for col in final_df.columns
]

测试效果 #

可以用下方模拟的原始数据测试，输出结果和预期格式完全一致：

# 模拟原始数据
raw_data = {
    'Year': [2018,2019,2020,2021,2022,2018,2019,2020],
    'ID': ['A','A','A','A','A','B','B','B'],
    'count_1': [1,2,3,4,5,10,20,30],
    'count_2': [10,20,30,40,50,100,200,300],
    'code': ['X001','X001','X001','X001','X001','Y002','Y002','Y002']
}
df = pd.DataFrame(raw_data)

运行上述代码后得到的final_df结构如下：

特殊场景适配 #

如果数据中存在同一ID对应多个code的情况，可以将code加入聚合逻辑，指定取第一个出现的code即可：

grouped_df = df.groupby(['ID', 'year_period'], as_index=False).agg(
    code = ('code', 'first'),
    count_1 = ('count_1', 'sum'),
    count_2 = ('count_2', 'sum')
)

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