переписать итерацию в .apply функцию с несколькими входами pandas

#python #pandas #dataframe #iteration #apply

Вопрос:

У меня есть два фрейма данных. Один содержит несколько электростанций и их соответствующее местоположение по долготе и широте, каждое в одном столбце. Другой фрейм данных содержит несколько подстанций, также с long и lat . То, что я пытаюсь сделать, это назначить электростанции ближайшим подстанциям.

 df1 = pd.DataFrame{'ID_pp':['p1','p2','p3','p4'],'x':[12.644881,11.563269, 12.644881,  8.153184], 'y':[48.099206, 48.020081, 48.099206, 49.153766]}
df2 = pd.DataFrame{'ID_ss':['s1','s2','s3','s4'],'x':[9.269, 9.390, 9.317, 10.061], 'y':[55.037, 54.940, 54.716, 54.349]}
 

Я нашел это решение, которое, по сути, именно то, что мне нужно:

 import pandas as pd
import geopy.distance



for i,row in df1.iterrows(): # A
    a = row.x, row.y
    distances = []
    for j,row2 in df2.iterrows(): # B
        b = row2.x, row2.y
        distances.append(geopy.distance.geodesic(a, b).km)

    min_distance = min(distances)
    min_index = distances.index(min_distance)

 
    df1['assigned_to'] =  min_index
 

это работает, но невероятно медленно, а мой набор данных довольно большой. Я думаю, что подход с использованием .apply был бы намного быстрее, но я не могу придумать, как его использовать. У кого-нибудь есть идея, как переписать вышеупомянутую функцию в подход .apply, который не требует итерации?

Ответ №1:

Решение в 8 раз быстрее (вдохновленное этой темой) :

 def closest_node(node, nodes):
    nodes = np.asarray(nodes)
    deltas = nodes - node
    dist_2 = np.einsum('ij,ij->i', deltas, deltas)
    return np.argmin(dist_2)

df2_array = df2[["x","y"]].values

df1['assigned_to'] = df1.apply(lambda x: closest_node2(np.array([x.x, x.y]),df2_array) ,axis=1)
 

Сравнение скорости :

перед:

 %%timeit
for i,row in df1.iterrows(): # A
    a = row.x, row.y
    distances = []
    for j,row2 in df2.iterrows(): # B
        b = row2.x, row2.y
        distances.append(geopy.distance.geodesic(a, b).km)

    min_distance = min(distances)
    min_index = distances.index(min_distance)

 
    df1['assigned_to'] =  min_index
 

дает 8.75 ms ± 108 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

после :

 %%timeit
df2_array = df2[["x","y"]].values

df1['assigned_to'] = df1.apply(lambda x: closest_node2(np.array([x.x, x.y]),df2_array) ,axis=1)
 

дает 1.7 ms ± 31.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)