如何在Pandas中为first为green的行匹配后续首个second为green的时间

现有如下结构的DataFrame，包含time（时间列，示例为数字，实际为datetime对象或日期字符串）、first和second列：

t   first   second
1   grey    red
2   green   red
3   red     red
4   grey    green
5   green   red
6   grey    green
7   green   red
8   red     red

可通过以下代码创建该DataFrame：

import pandas as pd

dfx = pd.DataFrame(
    {
        'time': [1,2,3,4,5,6,7,8],
        'first': ['grey', 'green', 'red', 'grey', 'green', 'grey', 'green', 'red'], 
        'second': ['red', 'red', 'red', 'green', 'red', 'green', 'red', 'red']
    }
)

需求：筛选出first列值为green的行，并为这些行添加后续首个second列值等于green的时间，期望输出结果如下：

t   first   t_second
2   green   4
5   green   6
7   green   NaN

方法一：简单遍历实现（适合小数据集） #

先提取所有second为green的时间点，再为每个first为green的行匹配后续第一个符合条件的时间：

import pandas as pd

# 创建原始DataFrame
dfx = pd.DataFrame(
    {
        'time': [1,2,3,4,5,6,7,8],
        'first': ['grey', 'green', 'red', 'grey', 'green', 'grey', 'green', 'red'], 
        'second': ['red', 'red', 'red', 'green', 'red', 'green', 'red', 'red']
    }
)

# 提取所有second为green的时间点，按顺序排列
green_times = dfx[dfx['second'] == 'green']['time'].tolist()

# 筛选first为green的行作为结果基础
result = dfx[dfx['first'] == 'green'].copy()
result = result.rename(columns={'time': 't'})

# 为每个green行匹配后续第一个green的second时间
def find_next_green_time(current_time):
    candidates = [t for t in green_times if t > current_time]
    return candidates[0] if candidates else None

result['t_second'] = result['t'].apply(find_next_green_time)

# 重置索引使结果更整洁
result = result.reset_index(drop=True)

print(result)

运行输出：

t  first  t_second
0  2  green       4.0
1  5  green       6.0
2  7  green        NaN

方法二：向量化实现（适合大数据集） #

利用searchsorted的二分查找特性，效率远高于遍历，适合处理大规模数据：

import pandas as pd

dfx = pd.DataFrame(
    {
        'time': [1,2,3,4,5,6,7,8],
        'first': ['grey', 'green', 'red', 'grey', 'green', 'grey', 'green', 'red'], 
        'second': ['red', 'red', 'red', 'green', 'red', 'green', 'red', 'red']
    }
)

# 标记second为green的行并提取对应时间
green_mask = dfx['second'] == 'green'
green_times = dfx.loc[green_mask, 'time']

# 筛选first为green的行并整理格式
result = dfx[dfx['first'] == 'green'].rename(columns={'time': 't'}).reset_index(drop=True)

# 用二分查找找到每个t对应的第一个大于它的green_time位置
positions = green_times.searchsorted(result['t'], side='right')

# 匹配时间，超出范围的设为NaN
result['t_second'] = green_times.iloc[positions].reset_index(drop=True)
result.loc[positions >= len(green_times), 't_second'] = None

print(result)

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