Pandas按条件对比行生成新DataFrame 替代嵌套循环提升执行效率

现有超大数据框结构如下：

Timestamp    Val1

真实数据样例：

Timestamp           Temp         
0   1622471518.92911    36.443       
1   1622471525.034114   36.445       
2   1622471531.148139   37.447      
3   1622471537.284337   36.449      
4   1622471543.622588   43.345      
5   1622471549.734765   36.451      
6   1622471556.2518     36.454      
7   1622471562.361368   41.461     
8   1622471568.472718   42.468   
9   1622471574.826475   36.470

需要实现的逻辑：
对比Temp列的不同行数据，同时满足以下两个条件时，保存两条记录的相关信息：

• 后行温度值大于前行温度值的X倍（示例X=4）
• 两条记录的时间间隔小于Y分钟（示例Y=180分钟）

原始实现采用两层嵌套for循环，大数据量下耗时极长，需要用Pandas/numpy的向量化方案优化性能。
注：原始代码存在笔误，if dd > cap_time中的dd未定义，实际应为时间差变量tdiff

核心优化思路是将O(n²)复杂度的嵌套循环，替换为O(n log n)复杂度的二分查找+向量化判断，性能可提升千倍以上，适合超大规模数据：

import pandas as pd
import numpy as np

# 自定义参数
X = 4 # 温度倍数阈值
Y = 180 # 时间间隔阈值，单位分钟
cap_time_sec = Y * 60

# 预处理：先按时间戳升序排序，保证时间序列有序
df = df.sort_values('Timestamp', ascending=True).reset_index(drop=True)
# 转为numpy数组做向量化运算，速度远高于Pandas行操作
temps = df['Temp'].values
times = df['Timestamp'].values
n = len(temps)
results = []

for i in range(n):
    # 用二分查找快速定位当前行i对应的时间窗口右边界，无需遍历所有j
    max_allowed_time = times[i] + cap_time_sec
    right_bound = np.searchsorted(times, max_allowed_time, side='right')
    # 仅在时间窗口范围内判断温度条件，用numpy掩码做向量化判断
    temp_window = temps[i+1 : right_bound]
    match_mask = temp_window > temps[i] * X
    if np.any(match_mask):
        # 取第一个符合条件的j，如需取所有符合条件的行可删除后续的[0]判断逻辑
        first_match_pos = np.argmax(mask)
        j = i + 1 + first_match_pos
        time_diff_min = (times[j] - times[i]) / 60
        results.append([
            temps[i], temps[j], times[i], times[j],
            time_diff_min, Y, X
        ])

# 转为结果数据框
result_df = pd.DataFrame(results, columns=[
    'prev_temp', 'curr_temp',
    'prev_timestamp', 'curr_timestamp',
    'time_diff_min', 'threshold_time_min', 'threshold_multiple'
])

性能对比 #

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