如何高效在另一DataFrame中查找固定时间范围内的所有邻近记录

问题描述 #

现有两个DataFrame：order_df（包含订单ID和订单时间）和trade_df（包含交易ID和交易时间），需要为order_df的每一行，找出trade_df中交易时间处于订单时间前后指定分钟数（示例为3分钟）范围内的所有交易ID，最终将符合条件的交易ID合并为字符串列。

示例数据：

import datetime
import pandas as pd

order_dict = {
    'Order ID': ['Order 1', 'Order 2', 'Order 3', 'Order 4'],
    'Order Time': [datetime.datetime(2023, 3, 8, 7, 5, 0), datetime.datetime(2023, 3, 8, 7, 10, 0), datetime.datetime(2023, 3, 8, 7, 15, 0), datetime.datetime(2023, 3, 8, 7, 20, 0)]
}
trade_dict = {
    'Trade ID': ['Trade 1001', 'Trade 1002', 'Trade 1003', 'Trade 1004', 'Trade 1005'],
    'Trade Time': [datetime.datetime(2023, 3, 8, 7, 8, 0), datetime.datetime(2023, 3, 8, 7, 9, 0), datetime.datetime(2023, 3, 8, 7, 17, 0), datetime.datetime(2023, 3, 8, 7, 18, 0), datetime.datetime(2023, 3, 8, 7, 30, 0)]
}

order_df = pd.DataFrame(order_dict)
trade_df = pd.DataFrame(trade_dict)

预期输出（3分钟时间范围）：

Order ID           Order Time Nearest Trade ID
0   Order 1 2023-03-08 07:05:00        Trade 1001
1   Order 2 2023-03-08 07:10:00  Trade 1001, Trade 1002
2   Order 3 2023-03-08 07:15:00  Trade 1003, Trade 1004
3   Order 4 2023-03-08 07:20:00  Trade 1003, Trade 1004

原始低效实现 #

最初使用apply逐行遍历订单时间，筛选交易时间在范围内的记录，这种方法在数据量达10万条时耗时约3.5分钟，效率极低：

def fixed_distance_nearest_neighbors(x, fixed_distance, lookup_df, lookup_field, return_field):
    return ', '.join(lookup_df[lookup_df[lookup_field].between(x - datetime.timedelta(minutes=fixed_distance), x + datetime.timedelta(minutes=fixed_distance))][return_field])

order_df['Nearest Trade ID'] = order_df['Order Time'].apply(lambda x: fixed_distance_nearest_neighbors(x, 3, trade_df, 'Trade Time', 'Trade ID'))

高效实现方案 #

方法1：排序+二分查找（最优性能） #

核心思路：先对交易数据按时间排序，利用numpy.searchsorted对每个订单的时间范围进行二分查找，快速定位符合条件的交易起止索引，最后批量聚合结果。这种方法将逐行O(n)的查找转为O(log n)的二分查找，性能提升显著。

代码实现：

import numpy as np

# 预处理：对交易数据按时间排序，提取时间和ID的数组
trade_sorted = trade_df.sort_values('Trade Time').reset_index(drop=True)
trade_times = trade_sorted['Trade Time'].values
trade_ids = trade_sorted['Trade ID'].values

# 定义时间范围（此处为3分钟）
time_delta = datetime.timedelta(minutes=3)
order_times = order_df['Order Time'].values
lower_bounds = order_times - time_delta
upper_bounds = order_times + time_delta

# 用二分查找定位每个时间范围在交易数据中的起止位置
left_idx = np.searchsorted(trade_times, lower_bounds, side='left')
right_idx = np.searchsorted(trade_times, upper_bounds, side='right')

# 生成每个订单对应的交易ID字符串
order_df['Nearest Trade ID'] = [', '.join(trade_ids[left:right]) for left, right in zip(left_idx, right_idx)]

方法2：IntervalIndex匹配（简洁向量化） #

核心思路：为每个订单时间创建时间区间，利用pandas的IntervalIndex快速匹配交易时间所属的订单区间，再通过分组聚合得到结果，代码更简洁易读。

代码实现：

# 为每个订单时间创建包含前后3分钟的区间
time_delta = datetime.timedelta(minutes=3)
order_intervals = pd.IntervalIndex.from_arrays(
    order_df['Order Time'] - time_delta,
    order_df['Order Time'] + time_delta,
    closed='both'
)

# 对交易数据按时间排序，匹配每个交易所属的订单索引
trade_sorted = trade_df.sort_values('Trade Time')
matches = trade_sorted['Trade Time'].apply(lambda t: order_intervals.get_loc(t)).explode()

# 分组聚合交易ID，合并到原订单DataFrame
trade_groups = trade_sorted.loc[matches.index, 'Trade ID'].groupby(matches).agg(', '.join)
order_df['Nearest Trade ID'] = order_df.index.map(trade_groups).fillna('')

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