pandas 在另一个 Dataframe 中查找固定距离内的所有邻域

nx7onnlm  于 2023-03-11  发布在  其他
关注(0)|答案(2)|浏览(110)

我试图找出在另一个数据框中找到所有固定距离的邻居的最快方法。
为了更好地说明我的问题,我做了一个假设的例子,假设有两个 Dataframe order_dftrade_df,每个 Dataframe 有两列--一列是ID,另一列是datetime。

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中的每一行,我查找trade_df中的所有Trade ID,其中Trade TimeOrder Time之前或之后的任意分钟(例如,3分钟)内
例如,当固定距离= 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

我做了如下尝试:

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'))

但是我相信这并不是一个有效的方法,具体来说,当order_dftrade_df的大小增长到很大的数目时,上述方法是极其缓慢的,例如:

import random
import time

def generate_random_date(n):
    random_date = []
    for i in range(0, n):
        random_date.append(datetime.datetime(2023, 3, 8, random.randint(0, 23), random.randint(0, 59), random.randint(0, 59)))
    return random_date

n = 100000
order_dict = {
    'Order ID': ['Order ' + str(x) for x in range(1, n + 1)],
    'Order Time': generate_random_date(n)
}

trade_dict = {
    'Trade ID': ['Trade ' + str(x) for x in range(1, n + 1)],
    'Trade Time': generate_random_date(n)
}

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

#Output 214.34553575515747

执行需要3.5分钟,我相信还有一些更像Python的方法来提高速度,非常感谢。

rm5edbpk

rm5edbpk1#

以下代码使用了BallTree
您使用的10.000生成的订单/交易示例处理时间约为5秒。它之所以“慢”是因为许多交易/订单对冲突。如果日期范围更分散,并且只有少数交易与订单关联,则速度更快。
我只展示了玩具示例的示例。完全相同的代码用于测试更大的示例。
对于order_dicttrade_dict,我使用以下代码转换为UNIX时间;

order = pd.DataFrame(order_dict)
trade = pd.DataFrame(trade_dict)

order["unix_time"] = (order["Order Time"] - pd.Timestamp("1970-01-01")) // pd.Timedelta('1s')
trade["unix_time"] = (trade["Trade Time"] - pd.Timestamp("1970-01-01")) // pd.Timedelta('1s')

我们准备树的数据:

from sklearn.neighbors import BallTree
import numpy as np

order_times = np.expand_dims( order.unix_time.values, axis=1 )
trade_times = np.expand_dims( trade.unix_time.values, axis=1 )

......并创建和查询我们的树;

tree = BallTree(trade_times, leaf_size=15, metric="manhattan")
idx, distance_in_seconds = tree.query_radius( order_times, r=60*3, sort_results=True, return_distance=True)

以上操作是实际的交易/订单链接,对于10.000个交易/订单大约需要5秒。构建树也需要时间。所以如果你知道你要重复这个操作,如果不需要,尽量不要重建树。

  • 注意,我们使用Manhattan距离,表示交易可以远离订单的绝对时间(秒
  • 在query_radius()中,我们使用sort_results=True, return_distance=True来确保输出是有序的,即最近的交易/订单排在第一位。
  • 请注意,radius60*3秒。
order["Nearest Trades"] = idx

会给予我们

Order ID          Order Time   unix_time Nearest Trades
0  Order 1 2023-03-08 07:05:00  1678259100            [0]
1  Order 2 2023-03-08 07:10:00  1678259400         [1, 0]
2  Order 3 2023-03-08 07:15:00  1678259700         [2, 3]
3  Order 4 2023-03-08 07:20:00  1678260000         [3, 2]

注意,对于Order 2,最近的交易是trade 1,就在trade 0之后,对于较大的示例,获取准确的名称会变得很麻烦,所以我将使用indici而不是名称。
还有distance_in_seconds,它返回以秒为单位的距离,它的形状与idx相同,但它包含的不是索引,而是对应索引的以秒为单位的距离。

mf98qq94

mf98qq942#

您可以使用how='cross'作为merge的参数(如果您的 Dataframe 不是太大):

trades = (order_df.merge(trade_df, how='cross')
                  .loc[lambda x: x['Order Time'].sub(x['Trade Time']).abs().lt('3min')]
                  .groupby('Order ID')['Trade ID'].agg(', '.join))

order_df['Nearest Trade'] = order_df['Order ID'].map(trades).fillna('')

输出:

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

相关问题