Pandas筛选连续4次及以上True值的分组记录及优化实现

初始测试数据 #

构造测试DataFrame的代码如下：

import pandas as pd
d = {'id': [1,1,1,1,1,1,2,2,2,2,3], 
     'log_date' : ['2021-01-05 07:23:00', '2021-01-05 07:24:00', '2021-01-05 07:25:00', '2021-01-05 07:26:00', '2021-01-05 07:27:00', '2021-01-05 07:28:00', '2021-01-06 07:23:00', '2021-01-06 07:24:00', '2021-01-06 07:25:00', '2021-01-06 07:23:00', '2021-01-05 07:20:00'], 
     'is_bool': [False,True,True,True,True,True,False,True,True,True,True ]}
df = pd.DataFrame(data=d)

需求目标 #

筛选出存在4个及以上连续True值的id，返回对应id、日志日期（取log_date的日期部分）、该id下最长连续True的长度，预期输出：

现有尝试 #

目前通过逐行for循环新建counter列，累加统计连续True的长度，伪代码如下：

for i in range(len(df)):
  if i == 0:
    df.loc[i,'counter'] = 1 if df.loc[i,'is_bool'] else 0
  elif df.loc[i,'id'] == df.loc[i-1,'id'] and df.loc[i,'is_bool']:
    df.loc[i,'counter'] = df.loc[i-1,'counter'] + 1
  else:
    df.loc[i,'counter'] = 0

待解决疑问：

• 如何使用apply方法实现上述计数逻辑？
• 是否存在更优的实现方式？

不推荐用apply实现这类逻辑。apply在处理依赖前序行计算结果的场景时，本质还是逐行串行遍历，和手写for循环的性能没有本质差异，代码可读性反而更低。
pandas提供了分组向量化的连续值统计方案，性能远高于循环和apply，实现步骤如下：

1. 将log_date转为日期时间类型，方便后续提取日期
2. 按id分组，对组内的is_bool列标记连续值块：连续相同的布尔值会被归为同一个块
3. 统计所有取值为True的连续块的长度
4. 按id聚合取最长连续True的长度，筛选长度≥4的id
5. 关联对应日期字段，调整列顺序得到最终结果

完整可运行代码：

# 转换日期格式
df['log_date'] = pd.to_datetime(df['log_date'])

# 标记每个id下的连续值块
df['block_id'] = df.groupby('id')['is_bool'].apply(
    lambda x: (x != x.shift()).cumsum()
).reset_index(drop=True)

# 统计每个True连续块的长度
true_block_stats = df[df['is_bool']].groupby(['id', 'block_id']).size().reset_index(name='consec_len')

# 取每个id下的最长连续True长度，筛选符合>=4要求的id
result = true_block_stats.groupby('id')['consec_len'].max().reset_index(name='max_count')
result = result[result['max_count'] >= 4]

# 关联对应日期，调整输出格式
date_map = df.assign(log_date=df['log_date'].dt.date).drop_duplicates('id')[['id', 'log_date']]
result = result.merge(date_map, on='id')[['id', 'log_date', 'max_count']]

运行后输出结果和预期完全一致。

如果数据量特别大，还可以用numba加速循环逻辑，但上述分组向量化方案是pandas原生实现中通用性和性能平衡最优的选择。

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