使用Rolling Window识别DataFrame中连续9个1并生成新列

问题背景 #

已有一个包含a、b列的DataFrame，已通过df['y'] = (df['a'] > df['b']).astype(int)生成了由0和1组成的y列。需要找出所有属于连续9个及以上1的行，将对应a列的值复制到新列z，不满足条件的行填充NaN。

实现代码 #

方法一：基于分组统计（直观高效）

import pandas as pd
import numpy as np

# 示例数据（替换为你的实际DataFrame即可）
data = {
    'a': [-0.143, -0.109, -0.118, -0.118, -0.090, -0.090, -0.075, -0.075, -0.065, -0.065, -0.042, -0.042, -0.064, -0.064, -0.055, -0.055, -0.068, -0.068, -0.074, -0.063, -0.074, -0.063, -0.077, -0.077, -0.082, -0.127, -0.095, -0.095, -0.108],
    'b': [-0.109, -0.108, -0.108, -0.113, -0.110, -0.108, -0.050, -0.059, -0.056, -0.053, -0.040, -0.039, -0.060, -0.057, -0.054, -0.053, -0.069, -0.056, -0.075, -0.076, -0.056, -0.069, -0.075, -0.050, -0.058, -0.056, -0.100, -0.094, -0.096]
}
df = pd.DataFrame(data)
df['y'] = (df['a'] > df['b']).astype(int)

# 1. 对连续的相同y值分组
df['group'] = (df['y'] != df['y'].shift()).cumsum()
# 2. 计算每个分组的行数
group_length = df.groupby('group')['y'].transform('count')
# 3. 标记出属于连续9个及以上1的行
df['is_valid'] = (df['y'] == 1) & (group_length >= 9)
# 4. 生成z列：满足条件取a的值，否则填充NaN
df['z'] = np.where(df['is_valid'], df['a'], np.nan)

# 可选：删除中间辅助列
df.drop(['group', 'is_valid'], axis=1, inplace=True)

方法二：基于滚动窗口（贴合需求中的rolling window要求）

import pandas as pd
import numpy as np

# 构造示例DataFrame（省略重复代码）
df['y'] = (df['a'] > df['b']).astype(int)

window_size = 9
# 1. 找出连续9个1的窗口结束位置（窗口和为9说明全是1）
df['is_end_of_sequence'] = df['y'].rolling(window=window_size).sum() == window_size
# 2. 将标记扩展到窗口内的所有行
for i in range(window_size):
    df['is_end_of_sequence'] = df['is_end_of_sequence'] | df['is_end_of_sequence'].shift(i, fill_value=False)
# 3. 生成z列
df['z'] = np.where(df['is_end_of_sequence'], df['a'], np.nan)

# 可选：删除辅助列
df.drop('is_end_of_sequence', axis=1, inplace=True)

代码说明 #

• 方法一：通过分组连续相同值，直接统计每组长度，筛选出长度≥9的1序列，逻辑清晰且效率更高，适合大数据集。
• 方法二：严格贴合滚动窗口的需求，先定位连续9个1的窗口终点，再将标记扩展到窗口内所有行，实现目标效果。

两种方法最终生成的z列都与示例数据一致：连续10个1的行对应z列取a的值，其余行填充NaN。

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