Pandas实现颜色单次触发标记（禁止前瞻）的技术问题

原代码实现 #

import numpy as np
import pandas_ta as ta
from pandas import DataFrame, pandas

df = pandas.DataFrame({"color": [None, None, 'blue', None, None, None, 'orange', None, None, None, None],
                       'bottom': [1, 2, 7, 5, 9, 9, 5, 4, 5, 5, 3],
                       'top': [5, 5, 11, 8, 10, 10, 9, 7, 10, 6, 7]})

print(df)

# lookback period
N = 3

# Pivot each color to own column and shift
df2 = (df.pivot(columns='color', values=['top', 'bottom'])
         .drop(columns=np.nan, level=1)
         .ffill(limit=N-1).shift()
       )

# compare current top with bottom & top from color occurance
out = df.join((df2['bottom'].le(df['top'], axis=0)
               & df2['top'].ge(df['top'], axis=0)).astype(int))
print(out)

原代码输出：

color  bottom  top  blue  orange
0     None       1    5     0       0
1     None       2    5     0       0
2     blue       7   11     0       0
3     None       5    8     1       0
4     None       9   10     1       0
5     None       9   10     1       0
6   orange       5    9     0       0
7     None       4    7     0       1
8     None       5   10     0       0
9     None       5    6     0       1
10    None       3    7     0       0

需求说明 #

每种颜色实例出现后，后续N=3行中仅能触发一次1（即每个颜色实例只能被"消耗"一次，连续出现多个同颜色实例则各自对应一次触发）。期望输出：

color  bottom  top  blue  orange
0     None       1    5     0       0
1     None       2    5     0       0
2     blue       7   11     0       0
3     None       5    8     1       0
4     None       9   10     0       0 --> blue已在第3行被消耗
5     None       9   10     0       0 --> blue已在第3行被消耗
6   orange       5    9     0       0
7     None       4    7     0       1
8     None       5   10     0       0
9     None       5    6     0       0 --> orange已在第7行被消耗
10    None       3    7     0       0

限制条件 #

禁止使用前瞻操作（如.shift(-3)、iloc[-1]等依赖未来行数据的方法）

解决方案 #

采用逐行遍历+状态维护的方式，完全基于历史和当前行数据处理，避免前瞻：

import numpy as np
import pandas as pd

df = pd.DataFrame({
    "color": [None, None, 'blue', None, None, None, 'orange', None, None, None, None],
    'bottom': [1, 2, 7, 5, 9, 9, 5, 4, 5, 5, 3],
    'top': [5, 5, 11, 8, 10, 10, 9, 7, 10, 6, 7]
})

N = 3

# 生成初始触发条件矩阵（和原逻辑一致）
df2 = (df.pivot(columns='color', values=['top', 'bottom'])
       .drop(columns=np.nan, level=1)
       .ffill(limit=N-1).shift()
       )
trigger = (df2['bottom'].le(df['top'], axis=0) & df2['top'].ge(df['top'], axis=0)).astype(int)

# 维护每个颜色的状态：剩余有效期（>0表示有未消耗的触发机会且在有效期内）
color_states = {col: 0 for col in trigger.columns}
result = trigger.copy()

for idx in range(len(df)):
    for col in trigger.columns:
        current_trigger = trigger.loc[idx, col]
        if color_states[col] > 0:
            if current_trigger == 1:
                # 触发成功，标记为1并消耗状态
                result.loc[idx, col] = 1
                color_states[col] = 0
            else:
                # 未触发，有效期减1
                result.loc[idx, col] = 0
                color_states[col] -= 1
        else:
            # 检查当前行是否为颜色出现行，重置有效期
            if df.loc[idx, 'color'] == col:
                color_states[col] = N
                result.loc[idx, col] = 0
            else:
                result.loc[idx, col] = 0

# 合并结果到原DataFrame
out = df.join(result)
print(out)

方案说明 #

1. 先通过原逻辑生成触发条件矩阵，判断每行是否满足对应颜色的触发要求
2. 用color_states字典记录每个颜色的剩余有效期：
   • 当颜色出现时，将该颜色的有效期设为N（后续N行可触发）
   • 逐行遍历，若颜色处于有效期内且触发条件满足，则标记为1并将有效期设为0（已消耗）
   • 若有效期内未触发，则有效期逐行递减，直到过期
3. 全程仅依赖当前行及之前的状态，完全符合禁止前瞻的要求

输出结果与期望一致：

color  bottom  top  blue  orange
0     None       1    5     0       0
1     None       2    5     0       0
2     blue       7   11     0       0
3     None       5    8     1       0
4     None       9   10     0       0
5     None       9   10     0       0
6   orange       5    9     0       0
7     None       4    7     0       1
8     None       5   10     0       0
9     None       5    6     0       0
10    None       3    7     0       0

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