基于Pandas筛选指定条件后的首行并创建新列的问题

原始数据与已有操作 #

首先定义初始DataFrame：

import pandas as pd
df = pd.DataFrame({'a': [20, 21, 333, 444, 1, 666], 'b': [20, 20, 20, 20, 20, 20], 'c': [222, 211, 2, 1, 100, 200]})

已实现的逻辑：找到第2个满足a >= b的行，为d列赋值'x'，代码如下：

mask = (df.a >= df.b)
df.loc[mask.cumsum().eq(2) & mask, 'd'] = 'x'

新增需求 #

需要找到同时满足以下条件的首行，为d列赋值'x'：

• 该行位于上述第一个掩码匹配的行之后
• 该行的c列数值大于b列数值

期望输出结果：

a   b    c    d
0   20  20  222    NaN
1   21  20  211    NaN
2  333  20    2    NaN
3  444  20    1    NaN
4    1  20  100    x
5  666  20  200    NaN

解决方案 #

方法一：分步定位索引 #

# 获取第2个满足a>=b的行的索引
target_idx = df.loc[(df.a >= df.b).cumsum().eq(2) & (df.a >= df.b)].index[0]

# 筛选目标索引之后且c>b的首行，赋值'x'
second_mask = (df.index > target_idx) & (df.c > df.b)
df.loc[second_mask.idxmax(), 'd'] = 'x'

方法二：链式索引定位 #

# 直接定位第一个掩码的目标行索引
first_target_idx = (df.a >= df.b).cumsum().eq(2).idxmax()

# 从目标行之后筛选c>b的首行索引
second_target_idx = df.loc[first_target_idx+1:].query('c > b').index[0]

# 赋值操作
df.loc[second_target_idx, 'd'] = 'x'

代码说明 #

• (df.a >= df.b).cumsum().eq(2).idxmax()：通过累加布尔值找到第2个满足a>=b的行的索引
• df.loc[first_target_idx+1:]：截取目标行之后的所有数据
• query('c > b')或(df.c > df.b)：筛选c大于b的行，idxmax()或.index[0]获取第一个符合条件的行索引

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