You need to enable JavaScript to run this app.
优惠活动
大模型
产品
解决方案
定价
更多

Pandas高效实现:为每行生成满足条件的最近历史索引

高效实现DataFrame每行匹配最近满足条件的前置索引

我需要为DataFrame中的每一行,返回之前最后一个满足lower列值≥当前行upper列值的索引。当前实现能得到预期结果,但不是真正的向量化操作,处理大型DataFrame时效率极低。

原实现代码

import pandas as pd

# 示例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['DATE'] = pd.to_datetime(df['DATE'])
df.set_index('DATE', inplace=True)

# 新增列存储符合条件的最近前置索引
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) 

print(df)

问题:如何重写代码实现最高效性能?


编辑补充:性能对比结果

感谢各位回复,在四种可行方案的性能对比中,Andrej Kesely提出的bisect方案表现最优。pyjanitor方案在数据量接近我的数据集时会触发内存分配错误,已排除。

性能测试结果

  • baseline(原方案):1min 35s ± 5.15 s per loop(2次运行,每次2循环的均值±标准差)
  • bisect方案:1.76 s ± 82.5 ms per loop(2次运行,每次2循环的均值±标准差)
  • enumerate方案:1min 13s ± 2.17 s per loop(2次运行,每次2循环的均值±标准差)

测试代码

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

def get_sample_df(rows=100_000):
    # 生成测试用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['DATE'] = pd.to_datetime(df['DATE'])
    df.set_index('DATE', inplace=True)

    return df


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


def get_pyjanitor():

    df = get_sample_df()
    df.reset_index(inplace=True)

    # 处理索引匹配
    left_df = df.assign(index_prev=df.index)
    right_df = df.assign(index_next=df.index)
    out=(left_df
        .conditional_join(
            right_df, 
            ('lower','upper','>='), 
            ('index_prev','index_next','<'), 
            df_columns='index_prev', 
            right_columns=['index_next','lower','upper'])
        )
    # 筛选最接近当前行的匹配项
    closest=out.index_next-out.index_prev
    grouper=[out.index_next, out.lower,out.upper]
    min_closest=closest.groupby(grouper).transform('min')
    closest=closest==min_closest
    # 映射回DATE索引
    prev=out.loc[closest,'index_prev']
    prev=df.loc[prev,'DATE'].array # 避免索引对齐问题
    index_next=out.loc[closest,'index_next']
    prev=pd.Series(prev,index=index_next)
    df = df.assign(prev=prev)
    return df


def get_bisect():
    df = get_sample_df()

    def get_prev_bs(lower, upper, _date):
        uniq_lower = sorted(set(lower))
        last_seen = {}

        for l, u, d in zip(lower, upper, _date):
            # 找到第一个≥u的元素索引
            idx = bisect_left(uniq_lower, u)

            max_date = None
            for lv in uniq_lower[idx:]:
                if lv in last_seen:
                    if max_date is None:
                        max_date = last_seen[lv]
                    elif last_seen[lv] > max_date:
                        max_date = last_seen[lv]
            yield max_date
            last_seen[l] = d

    df["prev"] = list(get_prev_bs(df["lower"], df["upper"], df.index))
    return df

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:
            if (any(j >= y for j in lower_list[0:i])):
                # 倒序查找最近的符合条件的索引
                for ll,dl in zip(reversed(lower_list[0:i]),reversed(date_list[0:i])):
                    if ll>=y:
                        new_list.append(dl)
                        break
                    else:
                        continue
            else:
                new_list.append(None)
    df['prev']=new_list
    df['prev']=pd.to_datetime(df['prev'])
    return df

print("baseline:")
%timeit -n 2 -r 2 get_baseline()

# 内存分配错误:Unable to allocate 37.2 GiB for an array with shape (4994299505,) and data type int64
# print("pyjanitor:")
# %timeit -n 2 get_pyjanitor()

print("bisect:")
%timeit -n 2 -r 2 get_bisect()

print("enumerate:")
%timeit -n 2 -r 2 get_enumerate()

内容的提问来源于stack exchange,提问作者TheBob

相关产品推荐
方舟 Agent Plan

超全模态模型 × Harness 升级,最新支持 Deepseek-V4.1-Flash、GLM-5.3 系列、Doubao-Seedream-5.0-pro、Kimi-K3 (部分), 限时 9.9 元起

最近更新时间:2026.07.03 14:35:02