Pandas无需循环实现两布尔列间条件滚动计数的方法

你要的效果可以通过pandas的分组累加向量化操作实现，完全不需要逐行for循环，百万级数据也可以快速跑完，代码如下：

import pandas as pd
import numpy as np

# 示例数据构造，你可以替换成自己的数据源
data = {
    'A': [False, True, False, True, False, False, True, False, False, False, False],
    'B': [False, False, False, False, True, False, False, True, True, False, True]
}
df = pd.DataFrame(data)

# 核心逻辑
## 1. 标记每行的触发事件类型：A列True记为A，B列True记为B，都为False记为空
event = pd.Series(np.where(df['A'], 'A', np.where(df['B'], 'B', pd.NA)), dtype='string')

## 2. 计算分组ID：每当触发的列变化时，分组ID+1
non_na_event = event.dropna()
group_change = non_na_event != non_na_event.shift()
group_id = pd.Series([pd.NA] * len(df), dtype='Int64')
group_id.loc[non_na_event.index] = group_change.cumsum()
# 前向填充所有行的分组ID，开头无触发事件的行分组统一为0
group_id = group_id.ffill().fillna(0).astype(int)

## 3. 每个分组内按触发列累加True的数量，得到count列
event_type = event.groupby(group_id).transform('first')
count_calc = (df['A'] * (event_type == 'A') + df['B'] * (event_type == 'B'))
df['count'] = count_calc.groupby(group_id).cumsum()

运行后得到的count列和你给出的样例完全一致。

逻辑说明 #

• 首先识别每一行的触发列是A还是B，过滤掉两列都为False的无触发行
• 每当触发列从A变成B或者从B变成A时，就生成一个新的分组，所有行的分组ID会通过前向填充补全
• 每个分组内只统计首次触发列的True值累计数量，直接用分组累加即可得到最终的count结果

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