Pandas高效查找历史条件匹配的最新索引：Bisect方法详解

import pandas as pd
import numpy as np

# 示例DataFrame
data = {'lower': [7, 1, 6, 1, 1, 1, 1, 11, 1, 1],
        'upper': [2, 3, 4, 5, 6, 7, 8, 9, 10, 11]}
df = pd.DataFrame(data=data)
df['DATE'] = pd.date_range('2020-01-01', periods=len(data['lower']))
df.set_index('DATE', inplace=True)

print("原始DataFrame:")
print(df)

# 低效方案：使用 df.apply
def get_most_recent_index_baseline(row, dataframe):
    # 查找当前行之前的所有行
    # 注意：row.name - pd.Timedelta(minutes=1) 确保只考虑严格早于当前行的记录
    previous_indices = dataframe.loc[:row.name - pd.Timedelta(minutes=1)]  
    # 筛选满足条件的记录，并返回最新的索引
    recent_index = previous_indices[previous_indices['lower'] >= row['upper']].index.max()
    return recent_index

# 应用函数到每一行
# df['prev_baseline'] = df.apply(lambda row: get_most_recent_index_baseline(row, df), axis=1) 
# print("\n低效方案结果:")
# print(df)

from bisect import bisect_left

def get_prev_with_bisect(lower_series, upper_series, date_index):
    """
    使用bisect模块高效查找满足条件的历史最新索引。

    参数:
    lower_series (pd.Series): DataFrame的'lower'列。
    upper_series (pd.Series): DataFrame的'upper'列。
    date_index (pd.DatetimeIndex): DataFrame的时间索引。

    返回:
    list: 包含每行对应的历史最新索引的列表。
    """
    # 获取所有不重复的lower值并排序，用于二分查找
    uniq_lower = sorted(set(lower_series))
    # last_seen 字典用于存储每个lower值最近出现的日期
    # 键为lower值，值为对应的最新日期
    last_seen = {}

    results = []

    # 遍历每一行数据
    for l, u, d in zip(lower_series, upper_series, date_index):
        max_date = None

        # 使用bisect_left查找在uniq_lower中，第一个大于或等于当前u的元素的索引
        # 这意味着从idx开始的所有uniq_lower元素都满足 lower >= u 的条件
        idx = bisect_left(uniq_lower, u)

        # 遍历所有满足条件的lower值
        for lv in uniq_lower[idx:]:
            if lv in last_seen:
                # 如果该lower值在之前出现过
                if max_date is None:
                    max_date = last_seen[lv]
                elif last_seen[lv] > max_date:
                    # 更新为更近的日期
                    max_date = last_seen[lv]

        results.append(max_date)

        # 更新last_seen字典：记录当前l值对应的最新日期d
        last_seen[l] = d

    return results

# 应用优化后的函数
df['prev_bisect'] = get_prev_with_bisect(df["lower"], df["upper"], df.index)

print("\nBisect方案结果:")
print(df)

import pandas as pd
import numpy as np
from bisect import bisect_left
import time

def get_sample_df(rows=100_000):
    # Sample DataFrame
    data = {'lower': np.random.default_rng(seed=1).uniform(1,100,rows),
            'upper': np.random.default_rng(seed=2).uniform(1,100,rows)}

    df = pd.DataFrame(data=data)
    df = df.astype(int)

    df['DATE'] = pd.date_range('2020-01-01', periods=len(data['lower']), freq="min")
    df.set_index('DATE', inplace=True)
    return df

# 基线方法 (get_baseline) - 与 get_most_recent_index_baseline 相同
def get_baseline():
    df = get_sample_df()
    def get_most_recent_index(row):
        previous_indices = df.loc[:row.name - pd.Timedelta(minutes=1)]  
        recent_index = previous_indices[previous_indices['lower'] >= row['upper']].index.max()
        return recent_index
    df['prev'] = df.apply(get_most_recent_index, axis=1) 
    return df

# Bisect 方法 (get_bisect) - 与 get_prev_with_bisect 相同
def get_bisect():
    df = get_sample_df()
    df["prev"] = get_prev_with_bisect(df["lower"], df["upper"], df.index)
    return df

# 朴素的enumerate循环方法 (get_enumerate)
def get_enumerate():
    df = get_sample_df()
    df.reset_index(inplace=True) # 重置索引方便列表操作

    date_list=df["DATE"].values.tolist()
    lower_list=df["lower"].values.tolist()
    upper_list=df["upper"].values.tolist()
    new_list=[]
    for i,(x,y) in enumerate(zip(lower_list,upper_list)):
        if i==0:
            new_list.append(None)
        else:
            found_date = None
            # 从后向前遍历，找到第一个满足条件的日期
            for ll,dl in zip(reversed(lower_list[0:i]),reversed(date_list[0:i])):
                if ll>=y:
                    found_date = dl
                    break
            new_list.append(found_date)
    df['prev']=new_list
    df['prev']=pd.to_datetime(df['prev'])
    return df

print("--- 性能测试 (100,000 行) ---")

start_time = time.time()
get_baseline()
print(f"Baseline (df.apply): {time.time() - start_time:.2f} seconds")

start_time = time.time()
get_bisect()
print(f"Bisect: {time.time() - start_time:.2f} seconds")

start_time = time.time()
get_enumerate()
print(f"Enumerate (Python loop): {time.time() - start_time:.2f} seconds")