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
相关产品推荐
相关产品推荐

