Python中基于列值用条件while循环实现数据行滚动9个月求和

实现思路 #

• 先按 Code 分组，每组内按 Dates 排序，确保数据时序逻辑正确
• 对每组内的每条数据，筛选出当前行日期及之后9个月内的所有行，计算这些行的 Value 总和

具体代码实现 #

使用 pandas 库处理，步骤如下：

1. 准备示例数据 #

import pandas as pd
from datetime import timedelta

# 构造符合需求的示例数据
data = {
    'Code': ['123', '123', '123', '123', '456', '456'],
    'Dates': pd.to_datetime(['2023-01-01', '2023-04-01', '2023-07-01', '2024-01-01', '2023-02-01', '2023-10-01']),
    'Value': [1, 2, 3, 4, 5, 6]
}
df = pd.DataFrame(data)

2. 基础分组计算逻辑 #

# 先按Code分组，每组内按Dates排序并重置索引
df_sorted = df.sort_values(['Code', 'Dates']).reset_index(drop=True)

# 定义每组内的计算函数
def calculate_9month_sum(group):
    sum_list = []
    for idx, row in group.iterrows():
        # 计算当前日期后推9个月的截止日期
        end_date = row['Dates'] + pd.DateOffset(months=9)
        # 筛选当前行及之后、日期不超过截止日的行
        mask = (group['Dates'] >= row['Dates']) & (group['Dates'] <= end_date)
        # 求和并加入结果列表
        total = group.loc[mask, 'Value'].sum()
        sum_list.append(total)
    # 新增列存储结果
    group['9MonthSum'] = sum_list
    return group

# 将函数应用到每个分组
result_df = df_sorted.groupby('Code', group_keys=False).apply(calculate_9month_sum)

3. 大数据量优化方案 #

如果数据集较大，循环遍历效率较低，可改用时间窗口滚动求和：

# 按Code分组后，将Dates设为索引
df_sorted = df_sorted.set_index('Dates')

# 使用9个月时间窗口的滚动求和，closed='left'确保包含当前行
rolling_sum = df_sorted.groupby('Code')['Value'].rolling('9M', closed='left').sum().reset_index()
# 重命名列并合并原数据
rolling_sum = rolling_sum.rename(columns={'Value': '9MonthSum'})
result_df = pd.merge(df_sorted.reset_index(), rolling_sum, on=['Code', 'Dates'], how='left')

结果说明 #

运行后 result_df 中的 9MonthSum 列即为需求结果：

• Code为123的第一行（2023-01-01），9个月截止日为2023-10-01，包含前3行，总和为1+2+3=6
• Code为123的第四行（2024-01-01），9个月截止日为2024-10-01，仅包含自身，总和为4
• Code为456的第一行（2023-02-01），9个月截止日为2023-11-01，包含两行，总和为5+6=11

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