如何用向量化方法为Pandas DataFrame按间隔条件设置列值

原始DataFrame #

import pandas as pd

df = pd.DataFrame({
    'Condition': ['A', 'B', 'A', 'A', 'A', 'B', 'A', 'C', 'A', 'A']
})

期望输出DataFrame #

expected_df = pd.DataFrame({
    'Condition': ['A', 'B', 'A', 'A', 'A', 'B', 'A', 'C', 'A', 'A'],
    'Output': [True, None, None, None, None, None, True, None, None, None]
})

规则说明 #

• 当df.Condition == 'A'时，df.Output = True；
• 在某行Output设为True后，连续3行无论Condition是否为'A'，df.Output都设为None。

可以完全通过向量化方法实现，无需循环。核心思路是：

1. 先标记所有满足Condition == 'A'的触发行；
2. 标记出触发行之后的连续3行（这些行需要强制设为None）；
3. 结合两个标记，仅在触发行且不在屏蔽范围内时设置True，其余设为None。

代码实现 #

import numpy as np

# 标记所有触发条件的行
trigger = df['Condition'] == 'A'

# 生成屏蔽标记：标记触发行之后的3行
# 卷积核[0,1,1,1]表示：触发行的下1、2、3行需要被屏蔽
mask = np.convolve(trigger.astype(int), [0, 1, 1, 1], mode='same') > 0

# 生成Output列
df['Output'] = np.where(trigger & ~mask, True, None)

验证结果 #

运行上述代码后，df的Output列将与期望输出完全一致。

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