Выборка pandas DF для соответствия второму DF в пределах ошибки

Question

Предположим, у меня есть два DF, скажем df1,df2 следующим образом:

import pandas as pd
import numpy as np
df1 = pd.DataFrame([[0,1,100],[1,1.1,120],[2,0.8,102]],columns=['id','a','b'])
df2 = pd.DataFrame([[0,0.5,110],[1,1.05,94],[2,0.96,145],[3,0.86,112],[4,1.3,97]],
                    columns=['id','a','b'])

print(df1)
   id    a    b
0   0  1.0  100
1   1  1.1  120
2   2  0.8  102
print(df2)
   id     a    b
0   0  0.50  110
1   1  1.05   94
2   2  0.96  145
3   3  0.86  112
4   4  1.30   97

Теперь предположим, что я выбираю некоторый размер интервала da,db. Я хочу для каждой строки в df1 выбрать случайную строку из df2, например abs(a1-a2)<da,abs(b1-b2)<db. То, что я сейчас делаю, - это очень грубая сила:

da = 0.2
db = 25
df2_list=[]
nbad = 0
for rid,row in df1.iterrows():
    ca = row['a']
    cb = row['b']
    c_df2 = df2[np.abs(df2['a']-ca)<da]\
            [np.abs(df2['b']-cb)<db]
    if len(c_df2) == 0:
        nbad+=1
        continue
    c_df2 = c_df2.sample()
    df2_list.append(c_df2['id'].values[0])

    
matched_df = df2[df2['id'].isin(df2_list)]
print(matched_df)

   id     a    b
1   1  1.05   94
3   3  0.86  112
4   4  1.30   97

Однако для моей реальной цели, где мой DF действительно большой, это очень медленно. Есть ли более быстрый способ достичь этого результата?

Answer 1

Вот решение:

da = 0.2
db = 25

res = pd.merge(df1.assign(dummy = 1), df2.assign(dummy = 1), on = "dummy").drop("dummy", axis = 1)
res = res[(np.abs(res.a_x - res.a_y) < da) & (np.abs(res.b_x - res.b_y) < db)]
res = res.groupby("id_x").apply(lambda x: x.sample(1))[["id_y", "a_y", "b_y"]]
res.index = res.index.droplevel(1)
print(res)

Результат:

      id_y   a_y  b_y
id_x                 
0        1  1.05   94
1        4  1.30   97
2        3  0.86  112