基于Pandas实现按ID分组且满足条件时重置并保持0的累计计数

需要在Pandas DataFrame中新增count列，规则如下：

• 当boolean字段为True时，连续累计计数；为False时重置计数并保持为0
• 按ID列分组，切换新ID时计数需重置
• 处理大数据集，禁止使用循环

预期输出示例：

ID  Boolean  Count
1   True     1
1   True     2
1   True     3
1   True     4
1   True     5
1   False    0
1   False    0 
1   False    0
1   False    0
1   True     1
1   True     2
1   True     3
2   True     1
2   True     2
2   True     3
2   True     4
2   False    0
2   False    0
2   False    0
2   True     1
2   True     2
2   True     3

解决方案 #

实现思路 #

1. 按ID分组，在组内标记计数重置的分界点：将连续的False段、以及每个True段的起始前的False作为独立分段
2. 按ID+分段标识分组，对Boolean列累计求和（True视为1，False视为0）
3. 最后将Boolean为False的行的Count值替换为0

代码实现 #

import pandas as pd

# 构造示例数据（可替换为你的数据集）
data = {
    'ID': [1]*12 + [2]*9,
    'Boolean': [True]*5 + [False]*4 + [True]*3 + [True]*4 + [False]*3 + [True]*3
}
df = pd.DataFrame(data)

# 步骤1：按ID分组，生成组内的分段重置标记
df['reset_tag'] = df.groupby('ID')['Boolean'].apply(lambda x: (~x).cumsum())

# 步骤2：按ID和reset_tag分组，计算组内累计计数
df['Count'] = df.groupby(['ID', 'reset_tag'])['Boolean'].cumsum()

# 步骤3：将Boolean为False的行的Count设为0
df['Count'] = df['Count'].where(df['Boolean'], 0)

# 清理中间列（可选）
df = df.drop('reset_tag', axis=1)

# 查看结果
print(df)

代码说明 #

• (~x).cumsum()：在每个ID组内，将False转换为True后累加，这样每遇到一个False都会生成新的分段标识，确保连续的True段被归为同一组
• groupby(['ID', 'reset_tag'])['Boolean'].cumsum()：在每个分段内对Boolean累加，实现连续True的计数
• where(df['Boolean'], 0)：快速将False对应的计数替换为0，避免循环处理

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