如何截取DataFrame至mask首次匹配位置？求最优实现方法

问题背景 #

给定如下DataFrame：

import pandas as pd
df = pd.DataFrame(
    {
        'a': [10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70],
        'b': [1, 1, 1, -1, -1, -2, -1, 2, 2, -2, -2, 1, -2],
    }
)

定义匹配规则的mask：

mask = (
    (df.b == -2) &
    (df.b.shift(1) > 0)
)

需求是截取df到mask首次匹配的索引之前的部分（mask首次匹配在索引9，预期输出到索引8），预期结果如下：

a  b
0  10  1
1  15  1
2  20  1
3  25 -1
4  30 -1
5  35 -2
6  40 -1
7  45  2
8  50  2

你自己实现的方法如下：

idx = df.loc[mask.cumsum().eq(1) & mask].index[0]
result = df.iloc[:idx]

最优实现方案 #

方案1：利用idxmax()直接获取首个匹配索引 #

布尔类型的Series中，True等价于1，False等价于0，idxmax()会直接返回第一个值为1（即True）的索引，代码简洁高效：

first_match_idx = mask.idxmax()
result = df.iloc[:first_match_idx]

如果需要处理mask无匹配的边界情况（避免索引错误），可以补充判断逻辑：

if mask.any():
    first_match_idx = mask.idxmax()
    result = df.iloc[:first_match_idx]
else:
    result = df.copy()  # 无匹配时返回原DataFrame，可按需调整

方案2：迭代匹配索引处理边界 #

这种方式更优雅地处理无匹配场景，当mask没有True值时，直接返回DataFrame的长度，截取整个数据集：

first_match_idx = next(iter(df[mask].index), len(df))
result = df.iloc[:first_match_idx]

方案对比 #

你的原方法是可行的，但通过mask.cumsum().eq(1) & mask定位首个匹配的写法相对繁琐。上面两种方案都能达到相同效果，且代码更简洁直观，时间复杂度和原方法一致（均为O(n)），属于更优的实现方式。

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