基于Pandas筛选掩码首行后满足区间条件的行并生成新列

问题背景 #

需为Pandas DataFrame生成新列x，核心逻辑基于指定掩码和区间匹配规则。

示例DataFrame #

import pandas as pd

df = pd.DataFrame(
    {
        'a': [100, 1123, 123, 100, 1, 0, 1],
        'b': [1000, 11123, 1123, 0, 55, 0, 1],
        'c': [100, 1123, 123, 999, 11, 50, 1],
        'd': [100, 1123, 123, 190, 1, 105, 1],
        'e': ['a', 'b', 'c', 'd', 'e', 'f', 'g'],
    }
)

目标输出 #

a      b     c     d  e   x
0   100   1000   100   100  a   NaN
1  1123  11123  1123  1123  b   NaN
2   123   1123   123   123  c   NaN
3   100      0   999   190  d   NaN
4     1     55    11     1  e   NaN
5     0      0    50   105  f   f
6     1      1     1     1  g   NaN

掩码定义 #

mask = (df.a > df.b)

实现步骤 #

• a) 找到满足mask条件的首行；
• b) 获取该行a列的值；
• c) 在mask首行之后的行中，找到第一个满足**该值处于c列与d列区间（包含等于）**的行；
• d) 获取该行e列的值，填充到新列x对应位置，其余行填充NaN。

示例逻辑说明 #

以提供的DataFrame为例：

• 满足mask的首行是行3；
• 该行a列值为100；
• 行3之后的行中，首个满足100处于c、d列区间的是行5，对应e列值为'f'；
• 因此行5的x列值为'f'，其余为NaN。

当前尝试代码 #

mask = (df.a > df.b)
val = df.loc[mask.cumsum().eq(1) & mask, 'a']

通用解决方案 #

以下方案支持所有边界场景（无满足掩码的行、掩码首行后无符合区间的行等）：

import pandas as pd

def generate_x_column(df):
    # 步骤a：定位满足掩码的首行
    mask = df['a'] > df['b']
    first_mask_idx = df[mask].index.min()
    
    # 无满足掩码的行时，x列全为NaN
    if pd.isna(first_mask_idx):
        df['x'] = pd.NA
        return df
    
    # 步骤b：获取目标值
    target_val = df.loc[first_mask_idx, 'a']
    
    # 步骤c：定位掩码首行后首个符合区间条件的行
    post_mask_rows = df.index > first_mask_idx
    interval_mask = (df['c'] <= target_val) & (target_val <= df['d'])
    target_row_idx = df[post_mask_rows & interval_mask].index.min()
    
    # 步骤d：填充x列
    df['x'] = pd.NA
    if not pd.isna(target_row_idx):
        df.loc[target_row_idx, 'x'] = df.loc[target_row_idx, 'e']
    
    return df

# 测试示例DataFrame
df = generate_x_column(df)
print(df)

测试用例 #

测试用例1：无满足掩码的行 #

df = pd.DataFrame(
    {'a': [100, 1123, 123, -1, 1, 0, 1], 
     'b': [1000, 11123, 1123, 0, 55, 0, 1],
     'c': [100, 1123, 123, 999, 11, 50, 1], 
     'd': [100, 1123, 123, 190, 1, 105, 1], 
     'e': ['a', 'b', 'c', 'd', 'e', 'f', 'g']}
)
df = generate_x_column(df)
print(df)

预期结果：x列全为NaN

测试用例2：掩码首行后无符合区间的行 #

df = pd.DataFrame(
    {'a': [100, 1123, 123, 100, 1, 0, 1], 
     'b': [1000, 11123, 1123, 0, 55, 0, 1], 
     'c': [100, 1123, 123, 999, 11, -1, 1], 
     'd': [100, 1123, 123, 190, 1, 10, 1], 
     'e': ['a', 'b', 'c', 'd', 'e', 'f', 'g']}
)
df = generate_x_column(df)
print(df)

预期结果：x列全为NaN

测试用例3：掩码首行在末尾，无后续行 #

df = pd.DataFrame(
    {'a': [100, 1123, 123, 1, 1, 0, 100], 
     'b': [1000, 11123, 1123, 0, 55, 0, 1], 
     'c': [100, 1123, 123, 999, 11, -1, 50], 
     'd': [100, 1123, 123, 190, 1, 10, 101], 
     'e': ['a', 'b', 'c', 'd', 'e', 'f', 'g']}
)
df = generate_x_column(df)
print(df)

预期结果：x列全为NaN

测试用例4：掩码首行后有多个符合条件的行，取第一个 #

df = pd.DataFrame(
    {'a': [100, 1123, 123, 100, 1, 1000, 1],
     'b': [1000, 11123, 1123, 0, 55, 0, 1],
     'c': [100, 1123, 123, 999, 11, 50, 500], 
     'd': [100, 1123, 123, 190, 1, 105, 2000], 
     'e': ['a', 'b', 'c', 'd', 'e', 'f', 'g']}
)
df = generate_x_column(df)
print(df)

预期结果：行5的x列值为'f'，其余为NaN

内容的提问来源于Stack Exchange，提问作者AmirX