如何在Pandas中根据start/end布尔条件实现区间计数？

规则 #

• 当start列的值为True时，开始计数
• 当end列的值为True时，停止计数并重置为0

输入数据 #

import pandas as pd
df = pd.DataFrame()
df['start'] = [False, True, False, False, False, True, False, False, False]
df['end'] = [False, False, False, True, False, False, False, True, False]

解决方案 #

通过分组累加的方式实现需求，代码如下：

# 标记每个计数周期的分组ID：遇到start或前一行是end时，开启新组
df['group'] = (df['start'] | df['end'].shift().fillna(False)).cumsum()

# 在每个组内从1开始累加计数
df['count'] = df.groupby('group').cumcount() + 1

# 将end为True的行计数重置为0
df['count'] = df.apply(lambda x: 0 if x['end'] else x['count'], axis=1)

# 对没有start触发的组（start前、end后的行）清零
df['count'] = df.where((df['group'].map(df.groupby('group')['start'].any())) & ~df['end'], 0)

# 输出符合预期的结果
print(df[['start', 'end', 'count']].rename(columns={'count': 'expected'}))

输出结果 #

start    end  expected
0  False  False         0
1   True  False         1
2  False  False         2
3  False   True         0
4  False  False         0
5   True  False         1
6  False  False         2
7  False   True         0
8  False  False         0

代码说明 #

• 分组标记：用(df['start'] | df['end'].shift().fillna(False)).cumsum()生成分组ID，每次碰到start或者上一行是end时，分组ID加1，把每个「start触发到end结束」的周期划进同一组。
• 组内计数：在每个组里用cumcount() + 1实现从1开始的累加。
• end行处理：直接将end为True的行计数设为0，符合停止并重置的规则。
• 无效周期清零：对没有start的分组（也就是未被start触发的行）统一清零，确保只有有效计数周期内才会产生非零数值。

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