Pandas中无循环查找third为True行后首满足second>first的行

问题回顾 #

原始DataFrame：

first   second   third
0     2       2      False
1     3       1      True
2     1       4      False
3     0       6      False
4     5       7      True
5     4       2      False
6     3       4      False
7     3       6      True

可通过以下代码创建：

import pandas as pd

df = pd.DataFrame(
    {
        'first': [2, 3, 1, 0, 5, 4, 3, 3], 
        'second': [2, 1, 4, 6, 7, 2, 4, 6], 
        'third': [False, True, False, False, True, False, False, True]
    }
)

需求：找出所有third列为True的行之后，首个满足second > first的行，优先避免循环实现。

期望输出：

first   second   third
2     1       4      False
6     3       4      False

无循环实现方法（使用merge_asof） #

利用Pandas的merge_asof函数可以高效完成这个需求，它支持按指定键进行向前匹配，正好符合“找后续首个符合条件行”的场景：

import pandas as pd

# 创建原始DataFrame
df = pd.DataFrame(
    {
        'first': [2, 3, 1, 0, 5, 4, 3, 3], 
        'second': [2, 1, 4, 6, 7, 2, 4, 6], 
        'third': [False, True, False, False, True, False, False, True]
    }
)

# 1. 筛选出所有third为True的行，保留原索引作为匹配键
df_true = df[df['third']].reset_index().rename(columns={'index': 'true_index'})

# 2. 筛选出所有满足second > first的行，保留原索引作为匹配键
df_valid = df[df['second'] > df['first']].reset_index().rename(columns={'index': 'valid_index'})

# 3. 使用merge_asof进行向前匹配：为每个True行找到后续首个有效行
merged = pd.merge_asof(
    df_true, 
    df_valid, 
    left_on='true_index', 
    right_on='valid_index', 
    direction='forward'
)

# 4. 提取目标行索引，去重后获取结果
target_indices = merged['valid_index'].dropna().unique()
result = df.loc[target_indices].sort_index()

print(result)

方法说明 #

• merge_asof的direction='forward'参数会为左表（df_true）的每一行，在右表（df_valid）中找到键值（索引）大于当前行键值的第一个记录，完美契合“后续首个”的要求。
• 最后对匹配到的索引去重、排序，再从原DataFrame中提取对应行，即可得到目标结果。

备选方法（含轻量循环） #

如果对merge_asof不熟悉，也可以结合bisect模块实现，循环仅作用于third为True的行索引，效率也较高：

import pandas as pd
import bisect

df = pd.DataFrame(
    {
        'first': [2, 3, 1, 0, 5, 4, 3, 3], 
        'second': [2, 1, 4, 6, 7, 2, 4, 6], 
        'third': [False, True, False, False, True, False, False, True]
    }
)

# 获取所有满足second > first的行索引
valid_indices = df.index[df['second'] > df['first']].tolist()

# 获取所有third为True的行索引
true_indices = df.index[df['third']].tolist()

# 逐个查找每个True行后续的首个有效行索引
target_indices = []
for idx in true_indices:
    pos = bisect.bisect_right(valid_indices, idx)
    if pos < len(valid_indices):
        target_indices.append(valid_indices[pos])

# 去重并提取结果
result = df.loc[sorted(set(target_indices))]
print(result)

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