正如评论中所暗示的,您需要一个更好的算法。您可以进行的一个优化是排序
col_B
,这样就可以执行二进制搜索来查找最近的元素。我将这两个实现放在一个
test.py
包含1000个随机元素的文件。下面是一些引导程序代码供您自己测试(
find_nearest2
"other_new_col"
).
O(N logn)解决方案保持与原始O(N^2)解决方案相同的行为:
import bisect
import numpy as np
import pandas as pd
# Original
def find_nearest1(row,array):
idx,val = min(enumerate(array['col_B']), key=lambda x: abs(x[1]-row['col_A']))
return array['col_C'][idx]
# Optimized
def find_nearest2(row,array):
idx = bisect.bisect_left(array['col_B'].values, row['col_A'])
arr_len = len(array['col_C'])
if idx == 0:
return array['col_C'].iloc[0]
elif idx == arr_len:
return array['col_C'].iloc[-1]
else:
diff1 = abs(array['col_B'].iloc[idx] - row['col_A'])
diff2 = abs(array['col_B'].iloc[idx-1] - row['col_A'])
m = min(diff1, diff2)
if np.isclose(m, diff1):
return array['col_C'].iloc[idx]
else:
return array['col_C'].iloc[idx-1]
np.random.seed(1) # Set seed for reproducability
size = 1000
df1_data = np.random.random(size)
df1_cols = ["col_A"]
df_1 = pd.DataFrame(df1_data, columns=df1_cols)
df2_data = {"col_B": np.random.random(size), "col_C": np.random.random(size)}
df_2 = pd.DataFrame.from_dict(df2_data)
if __name__ == "__main__":
# Original
df_1['new_col']=df_1.apply(lambda row: find_nearest1(row,df_2),axis=1)
# Optimized
df_2_sorted = df_2.sort_values('col_B')
df_1['other_new_col'] = df_1.apply(lambda row: find_nearest2(row, df_2_sorted), axis=1)
print(df_1)
这将输出(因此您可以验证是否保持了原始行为):
col_A new_col other_new_col
0 0.417022 0.842518 0.842518
1 0.720324 0.633461 0.633461
2 0.000114 0.327524 0.327524
3 0.302333 0.947542 0.947542
4 0.146756 0.985317 0.985317
5 0.092339 0.875530 0.875530
6 0.186260 0.348727 0.348727
7 0.345561 0.471819 0.471819
8 0.396767 0.674607 0.674607
9 0.538817 0.696878 0.696878
10 0.419195 0.695152 0.695152
...
990 0.812507 0.584883 0.584883
991 0.283802 0.497448 0.497448
992 0.527847 0.490089 0.490089
993 0.339417 0.905808 0.905808
994 0.554667 0.745292 0.745292
995 0.974403 0.615726 0.615726
996 0.311703 0.594144 0.594144
997 0.668797 0.571321 0.571321
998 0.325967 0.576978 0.576978
999 0.774477 0.963032 0.963032
一些性能结果:
bash-3.2$ ipython
Python 3.6.4 |Anaconda, Inc.| (default, Jan 16 2018, 12:04:33)
Type 'copyright', 'credits' or 'license' for more information
IPython 6.2.1 -- An enhanced Interactive Python. Type '?' for help.
In [1]: from test import find_nearest1, find_nearest2, df_1, df_2
In [2]: %timeit df_1['new_col']=df_1.apply(lambda row: find_nearest1(row,df_2),axis=1)
6.04 s ± 13.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [3]: %timeit df_2_sorted = df_2.sort_values('col_B')
201 µs ± 3.25 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [4]: df_2_sorted = df_2.sort_values('col_B')
In [5]: %timeit df_1.apply(lambda row: find_nearest2(row, df_2_sorted), axis=1)
111 ms ± 1.67 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
嗯。
