基于cond_a触发的b列100值3窗口滚动计数实现

问题背景 #

现有如下数据和预处理代码：

import pandas as pd

data = [100,100,100,100,10,100,10,100]
df = pd.DataFrame({'a':data, 'b':data})
df['cond_a'] = (df.a == 100).rolling(3).sum() == 3

处理后得到的DataFrame：

a    b  cond_a
0  100  100   False
1  100  100   False
2  100  100    True
3  100  100    True
4   10   10   False
5  100  100   False
6   10   10   False
7  100  100   False

需要添加cond_b_count列，要求：在首次观测到cond_a == True之后，从下一行开始，统计b列中100值在最多3个连续行（滚动窗口，窗口大小3，不足3行时按实际行数统计）内的出现次数，最终期望结果如下：

a    b  cond_a  cond_b_count
0  100  100   False           NaN
1  100  100   False           NaN
2  100  100    True           NaN
3  100  100    True           1.0
4   10   10   False           1.0
5  100  100   False           2.0
6   10   10   False           1.0
7  100  100   False           2.0

错误尝试 #

直接对全量数据应用滚动求和，未考虑触发条件的范围限制：

df['cond-b'] = (df.b == 100).rolling(3).sum()

正确解决方案 #

# 找到首次触发cond_a为True的索引
first_trigger_idx = df[df['cond_a']].index[0]

# 初始化目标列为空值
df['cond_b_count'] = pd.NA

# 截取触发点之后的子数据集，从触发点的下一行开始
post_trigger_data = df.loc[first_trigger_idx + 1:]

# 计算滚动窗口内100的出现次数：窗口大小3，不足3行时正常统计
rolling_counts = (post_trigger_data['b'] == 100).rolling(window=3, min_periods=1).sum()

# 将计算结果赋值回原DataFrame对应位置
df.loc[first_trigger_idx + 1:, 'cond_b_count'] = rolling_counts

代码说明 #

1. 定位触发点：找到第一个cond_a为True的索引，这是统计的起始边界。
2. 初始化空值：确保未触发区域的cond_b_count保持空值，避免无效数据干扰。
3. 截取目标数据：只对触发点之后的行进行统计，排除触发前的历史数据。
4. 滚动统计：使用rolling(window=3, min_periods=1)实现滚动窗口求和，min_periods=1保证窗口不足3行时仍能输出有效统计结果。
5. 赋值结果：将统计结果回填到原DataFrame的对应位置。

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