Pandas按日期分组求和：用历史stack超额填充缺口直至值为1

处理规则 #

• 第一步：按日期字段day_date分组，对stack字段做聚合求和
• 第二步：生成连续日频时间序列，原始数据中缺失的日期默认stack值为0
• 第三步：按日期先后逐天遍历：当日stack值超出1的部分记为超额量存入缓存，优先用缓存的超额量填充后续值为0的日期，每填充1天消耗1单位超额量，直到缓存耗尽
• 第四步：原始数据遍历完成后，若仍有剩余超额量，继续向后扩展日期序列，直到超额量全部消耗完毕，最终所有日期的stack值只能是0或1

注意：day_date为日频timestamp类型字段，下文示例为简化展示，用整数代表同一自然月内的日期。

示例场景 #

场景1：存在缺口、前期无累积超额栈 #

输入数据：

输出结果：

场景2：多段超额栈填充缺口 #

输入数据：

输出结果：

注：9日、10日stack值为1的原因是，6日的超额量已覆盖6-8日，8日剩余的2个单位超额量顺延填充了后续两日。

实现代码 #

import pandas as pd
from pandas.tseries.offsets import Day

def process_stack_dataframe(df: pd.DataFrame) -> pd.DataFrame:
    # 按日期分组聚合求和，按日期排序
    df = df.groupby('day_date', as_index=False)['stack'].sum()
    df = df.sort_values('day_date').reset_index(drop=True)
    
    surplus = 0  # 超额量缓存
    result = []
    current_date = df['day_date'].iloc[0]
    data_idx = 0
    max_original_date = df['day_date'].iloc[-1]

    while True:
        # 读取当日原始stack值，无原始记录则默认为0
        if data_idx < len(df) and df['day_date'].iloc[data_idx] == current_date:
            day_raw_stack = df['stack'].iloc[data_idx]
            data_idx += 1
        else:
            day_raw_stack = 0
        
        # 当日原始值全部加入超额缓存
        surplus += day_raw_stack
        # 当日值：缓存够1则取1，否则取0
        if surplus >= 1:
            today_value = 1
            surplus -= 1
        else:
            today_value = 0
            surplus = 0
        
        result.append({"day_date": current_date, "stack": today_value})

        # 终止条件：所有原始日期遍历完成，且超额缓存耗尽
        if current_date >= max_original_date and surplus <= 0:
            break
        
        # 日期向后推1天
        current_date += Day(1)
    
    return pd.DataFrame(result)


# 场景2测试用例
if __name__ == "__main__":
    test_df = pd.DataFrame({
        "day_date": pd.to_datetime([
            "2024-01-01", "2024-01-02", "2024-01-06", 
            "2024-01-08", "2024-01-15"
        ]),
        "stack": [0, 2, 3, 2, 1]
    })
    print(process_stack_dataframe(test_df))

代码运行后输出结果与上述场景示例完全匹配，传入真实timestamp格式的day_date字段即可直接使用。

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