Pandas按令牌累计消耗规则展开时间序列计算每日可用状态

核心规则 #

• 校验目标：判断参与者在对应日期是否至少持有1个可用令牌
• 计算逻辑：原始数据sum字段为当日新增令牌数量，参与者每日固定消耗1个令牌，历史累计结余的令牌可结转，用于覆盖后续无新增日期的消耗缺口
• 输出要求：
  • 生成连续日粒度时间序列，范围从数据集最早日期开始，到「累计令牌全部耗尽」或「数据集最晚记录日期」终止
  • 当日有可用令牌时sum字段记为1，无可用令牌记为0
  • 累计结余可覆盖无新增日期的缺口时，对应日期记为1；结余不足时缺口日期记为0；结余耗尽后再有新增令牌，重新开始累计结转

原有方案缺陷 #

之前采用pd.date_range按当日新增数直接生成日期范围再explode的实现，本质是将当日新增令牌全部绑定为新增日之后的连续占用，没有处理三类核心逻辑：

• 中间出现新增令牌时，结余需要动态累加
• 每日消耗后结余需要实时扣减，缺口日期按扣减后的结余判断是否可覆盖
• 序列末尾需要按实际剩余结余长度延伸，不能多算无效日期
  直接导致两个测试用例输出错误：
• 测试用例1：2021-01-08无结余应记为0，原有方案无法正确扣减跨日期的累计结余
• 测试用例2：累计令牌应覆盖至2021-01-06，原有方案错将1月6日记为0，还额外生成了2021-01-11的有效记录

正确实现代码 #

核心逻辑为逐日滚动计算令牌结余，避免提前展开日期导致的逻辑错位，代码可直接复用：

import pandas as pd

def get_daily_token_status(input_df):
    df = input_df.copy()
    # 基础预处理：日期格式转换，同日期多条新增记录求和
    df['date'] = pd.to_datetime(df['date'])
    df = df.groupby('date', as_index=False)['sum'].sum()

    # 第一步：生成数据集覆盖范围内的连续日期序列，无新增的日期填充0
    start_date = df['date'].min()
    data_end_date = df['date'].max()
    daily_df = pd.DataFrame({'date': pd.date_range(start_date, data_end_date, freq='D')})
    daily_df = daily_df.merge(df, on='date', how='left').fillna({'sum': 0})

    # 第二步：逐日滚动计算结余、当日可用状态
    balance = 0
    daily_status = []
    for _, row in daily_df.iterrows():
        balance += row['sum']
        if balance >= 1:
            # 有可用令牌，当日消耗1个
            daily_status.append(1)
            balance -= 1
        else:
            # 无可用令牌，结余不记负数
            daily_status.append(0)
            balance = 0
    daily_df['sum'] = daily_status

    # 第三步：处理末尾剩余结余，向后延伸到结余耗尽为止
    current_end = data_end_date
    while balance > 0:
        current_end += pd.Timedelta(days=1)
        daily_df = pd.concat([
            daily_df,
            pd.DataFrame({'date': [current_end], 'sum': [1]})
        ], ignore_index=True)
        balance -= 1

    return daily_df

测试用例验证 #

测试用例1 #

输入：

data1 = pd.DataFrame({
    'date':['2021-01-01','2021-01-02','2021-01-03','2021-01-06','2021-01-09','2021-01-10'],
    'sum':[2,2,2,1,2,1]
})
print(get_daily_token_status(data1))

输出符合预期：2021-01-01至2021-01-07均为1（累计结余覆盖1月4、5、7日缺口），2021-01-08记为0（结余耗尽），后续日期按新增规则正常计算。

测试用例2 #

输入：

data2 = pd.DataFrame({
    'date':['2021-01-01', '2021-01-02', '2021-01-10'], 
    'sum':[2,4,1]
})
print(get_daily_token_status(data2))

输出符合预期：2021-01-01至2021-01-06均为1（前两日累计新增6个令牌，覆盖6天消耗），2021-01-07至2021-01-09记为0（结余耗尽无新增），2021-01-10记为1，无多余的2021-01-11记录。

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