如何为Pandas DataFrame添加列x：查找首个超动态范围的open对应date

原始DataFrame #

import pandas as pd

df = pd.DataFrame(
    {
        'date': list('abcdefghi'),
        'open': [6, 8, 12, 5, 6, 22, 19, 1, 3],
        'high': [10, 12, 20, 5, 2, 44, 11, 12, 5],
        'low': [5, 7, 3, 1, 3, 18, 12, 1, 7],
    }
)

期望输出 #

date  open  high  low     x
0    a     6    10    5     c
1    b     8    12    7     NaN
2    c    12    20    3     f
3    d     5     5    1     NaN
4    e     6     2    3     NaN
5    f    22    44   18     h 
6    g    19    11   12     NaN
7    h     1    12    1     NaN
8    i     3     5    7     NaN

需求规则 #

• 初始范围为第一行的[low, high]即[5,10]
• 从当前范围的起始行之后，查找首个open值超出该范围（open > high 或 open < low）的行，将该行的date填入触发此次查找的行的x列
• 更新范围为找到的该行的[low, high]，重复上述步骤，直到找不到符合条件的行

用户尝试的代码 #

mask1 = (df.open > df.high)
mask2 = (df.open < df.low)
df.loc[mask1.cumsum().eq(1) | mask2.cumsum().eq(1), 'x'] = df.date

由于需求需要迭代更新范围并逐轮查找，无法通过一次性掩码实现，我们可以用循环迭代的方式处理：

import pandas as pd

df = pd.DataFrame(
    {
        'date': list('abcdefghi'),
        'open': [6, 8, 12, 5, 6, 22, 19, 1, 3],
        'high': [10, 12, 20, 5, 2, 44, 11, 12, 5],
        'low': [5, 7, 3, 1, 3, 18, 12, 1, 7],
    }
)

# 初始化x列为缺失值
df['x'] = pd.NA

current_idx = 0  # 初始起始行索引
while current_idx < len(df):
    # 获取当前行定义的范围
    current_low = df.loc[current_idx, 'low']
    current_high = df.loc[current_idx, 'high']
    
    # 从当前行的下一行开始，筛选open超出范围的行
    mask = (df['open'] > current_high) | (df['open'] < current_low)
    candidates = df.loc[current_idx+1:, :][mask]
    
    if not candidates.empty:
        # 取第一个符合条件的行的date，填入当前行的x列
        target_date = candidates.iloc[0]['date']
        df.loc[current_idx, 'x'] = target_date
        # 更新起始索引为找到的行，继续下一轮查找
        current_idx = candidates.index[0]
    else:
        # 无符合条件的行，退出循环
        break

print(df)

运行后输出与期望一致：

date  open  high  low     x
0    a     6    10    5     c
1    b     8    12    7  <NA>
2    c    12    20    3     f
3    d     5     5    1  <NA>
4    e     6     2    3  <NA>
5    f    22    44   18     h
6    g    19    11   12  <NA>
7    h     1    12    1  <NA>
8    i     3     5    7  <NA>

代码说明 #

1. 先初始化x列为缺失值，避免后续出现未赋值的情况
2. 用current_idx标记当前需要处理的行，从第一行开始
3. 每轮循环中，先获取当前行的范围，再从下一行开始查找首个超出范围的open值
4. 找到目标行后，将其date填入当前行的x列，并更新current_idx为目标行索引，继续下一轮
5. 当找不到符合条件的行时，终止循环

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