在 X 公里内找到最近的邻居以及他们的距离

Question

我有

df1

和

df2

以及 id 和纬度和经度坐标。

我当前的代码为 df1 中的每一行返回 df2 中的最近邻居及其相应的距离（以公里为单位）。

我想为 df1 中的每一行获取 X km 内 df2 中最近邻居的列表（一列用于他们的 ID，一列用于他们的相应距离）。

两个数据框包含很多点，这就是我使用pandarallel和KDTree的原因。

代码如下：

pandarallel.initialize(nb_workers=64, progress_bar=True)

# Convert latitude and longitude to radians
df1_df['lat_rad'] = df1_df['latitude'].apply(lambda x: x * (3.14159 / 180))
df1_df['lon_rad'] = df1_df['longitude'].apply(lambda x: x * (3.14159 / 180))

df2_df['lat_rad'] = df2_df['latitude'].apply(lambda x: x * (3.14159 / 180))
df2_df['lon_rad'] = df2_df['longitude'].apply(lambda x: x * (3.14159 / 180))

# Create KDTree using latitude and longitude columns
tree = KDTree(df2_df[['lat_rad', 'lon_rad']])

# Function to find nearest point in df2_df for each row in df1_df
def find_nearest(row):
    lat_rad = row['lat_rad']
    lon_rad = row['lon_rad']
    
    # Query the KDTree for the nearest point
    dist, idx = tree.query([(lat_rad, lon_rad)])
    
    # Get the corresponding df2_id and convert distance to km
    nearest_df2_id = df2_df.loc[idx, 'df2_id'].values[0]
    distance_km = dist[0] * 6371
    
    return nearest_df2_id, distance_km

# Apply the function to each row in df1_df using pandarallel
df1_df[['nearest_df2_id', 'distance_km']] = df1_df.parallel_apply(find_nearest, axis=1, result_type='expand')

# Remove the temporary columns
df1_df.drop(['lat_rad', 'lon_rad'], axis=1, inplace=True)

# Convert df2_id from float to int
df1_df.nearest_df2_id = df1_df.nearest_df2_id.astype('int')