分组计算满足条件的连续天数,条件不满足时重置累计值
问题描述
现有如下DataFrame数据(已按['op_id', 'si_id', 'date']排序):
import pandas as pd data = [['X','A','2022-05-01',True,True], ['X','A','2022-05-02',True,True], ['X','A','2022-05-03',True,False], ['X','A','2022-05-04',False,True], ['X','A','2022-05-05',False,True], ['X','A','2022-05-06',True,True], ['X','A','2022-05-07',True,False], ['X','B','2022-05-10',True,True], ['X','B','2022-05-11',True,True], ['X','B','2022-05-13',True,True], ['X','B','2022-05-14',True,True], ['X','B','2022-05-27',True,True], ['X','B','2022-05-28',False,True], ['X','B','2022-05-29',True,True], ['Y','C','2022-05-01',True,True], ['Y','C','2022-05-02',True,True], ['Y','C','2022-05-03',True,True], ['Y','C','2022-05-04',False,True], ['Y','C','2022-05-05',False,True], ['Y','C','2022-05-06',False,True]] columns = ['op_id','si_id','date','activity_level_1','activity_level_2'] df = pd.DataFrame(data, columns=columns).astype({'date':'datetime64[ns]'})
需求:按['op_id', 'si_id']分组,针对每个activity_level_n布尔列,计算值为True的连续天数,当值为False或日期不连续时重置累计值,最终得到类似如下的输出:
op_id si_id date activity_level_1 activity_level_2 cum_activity_level_1 cum_activity_level_2 0 X A 2022-05-01 True True 1 1 1 X A 2022-05-02 True True 2 2 2 X A 2022-05-03 True False 3 1 3 X A 2022-05-04 False True 1 2 4 X A 2022-05-05 False True 1 3 5 X A 2022-05-06 True True 1 4 6 X A 2022-05-07 True False 3 1 7 X B 2022-05-10 True True 1 1 8 X B 2022-05-11 True True 2 2 9 X B 2022-05-13 True True 1 1 10 X B 2022-05-14 True True 2 2 11 X B 2022-05-27 True True 1 1 12 X B 2022-05-28 False True 1 2 13 X B 2022-05-29 True True 1 3 14 Y C 2022-05-01 True True 1 1 15 Y C 2022-05-02 True True 2 2 16 Y C 2022-05-03 True True 3 3 17 Y C 2022-05-04 False True 1 4 18 Y C 2022-05-05 False True 1 5 19 Y C 2022-05-06 False True 1 6
尝试过以下代码,但结果不符合预期,累计值未正确重置:
day = pd.Timedelta('1d') breaks = (df['date'].diff() != day) | (~df['activity_level_1']) df['breaks'] = breaks df['cum_activity_level_1'] = df.activity_level_1.mask(breaks, False).groupby([df.activity_level_1, df.breaks]).cumsum()
需要找到高效的实现方式(数据集约700万行,优先避免apply)。
解决方案
核心思路是为每个分组内的每个activity列生成重置标记,再基于标记分组累计。以下是高效的向量化实现代码:
import pandas as pd # 1. 按op_id, si_id分组,计算日期是否连续 df['date_continuous'] = df.groupby(['op_id', 'si_id'])['date'].diff() == pd.Timedelta('1d') # 2. 提取所有需要处理的activity列 activity_cols = [col for col in df.columns if col.startswith('activity_level_')] # 3. 遍历每个activity列计算累计值 for col in activity_cols: # 生成重置标记:值为False 或 日期不连续时需要重置 reset_flag = ~df[col] | ~df['date_continuous'] # 生成分组键:累计重置标记的次数,每次重置则分组键递增 group_key = df.groupby(['op_id', 'si_id'])[reset_flag].cumsum() # 按分组键累计计数,从1开始 df[f'cum_{col}'] = df.groupby(['op_id', 'si_id', group_key]).cumcount() + 1 # 当activity值为False时,将累计值设为1(匹配示例输出) df.loc[~df[col], f'cum_{col}'] = 1 # 可选:删除中间生成的辅助列 df.drop('date_continuous', axis=1, inplace=True)
代码解释
- 日期连续标记:通过
groupby.diff()计算分组内的日期差,判断是否为1天,得到date_continuous列。 - 重置标记:
reset_flag为True时表示需要重置累计,覆盖了activity值为False、日期不连续两种场景。 - 分组键生成:用
cumsum()累计重置标记的次数,同一个连续有效序列的行拥有相同的分组键。 - 累计计数:
cumcount()+1在每个分组内生成从1开始的连续计数,最后修正activity为False的行的累计值为1。
该实现全程使用向量化操作,避免了apply,适合处理百万级别的大数据量。
内容的提问来源于stack exchange,提问作者filpa
相关产品推荐
相关产品推荐

