Python中高效从大型数据集查找最近地理位置的方案咨询

需求背景 #

需为含500,000条地址的数据集，每条从另一个1,000条的数据集里找出3个最近的地理位置（均含经纬度）。现有geopy+pandas逐行计算的方案效率不足，寻求最优实现，可使用免费GPU库。

现有实现代码 #

import pandas as pd
from geopy.distance import geodesic

# Sample data setup
data = {
    "Zip Code": ["10115", "20095", "50667", "80331", "70173", "40210", "41460", "45127", "47051", "40474"],
    "City": ["Berlin", "Hamburg", "Cologne", "Munich", "Stuttgart", "Düsseldorf", "Neuss", "Essen", "Duisburg", "Düsseldorf Airport"],
    "Latitude": [52.5323, 53.5503, 50.9367, 48.1372, 48.7833, 51.2217, 51.1981, 51.4556, 51.4344, 51.2895],
    "Longitude": [13.3846, 9.9930, 6.9540, 11.5755, 9.1815, 6.7763, 6.6913, 7.0116, 6.7623, 6.7668]
}
df = pd.DataFrame(data)

data_df2 = {
    "Address": ["my_location"],
    "ZipCode": ["40468"],
    "Latitude": [51.28472232436951],
    "Longitude": [6.7865234073914005]
}
df2 = pd.DataFrame(data_df2)

# Distance calculation
def calculate_distance(lat1, lon1, lat2, lon2):
    return geodesic((lat1, lon1), (lat2, lon2)).kilometers

def find_nearest_spots(address_lat, address_lon):
    distances = df.apply(lambda row: calculate_distance(address_lat, address_lon, row['Latitude'], row['Longitude']), axis=1)
    nearest_indices = distances.nsmallest(3).index
    nearest_spots = df.loc[nearest_indices]
    return pd.Series({
        'Nearest_1_Spot': nearest_spots.iloc[0]['City'],
        'Nearest_1_Dist': distances.iloc[nearest_indices[0]],
        'Nearest_2_Spot': nearest_spots.iloc[1]['City'],
        'Nearest_2_Dist': distances.iloc[nearest_indices[1]],
        'Nearest_3_Spot': nearest_spots.iloc[2]['City'],
        'Nearest_3_Dist': distances.iloc[nearest_indices[2]]
    })

df2[['Nearest_1_Spot', 'Nearest_1_Dist', 'Nearest_2_Spot', 'Nearest_2_Dist', 'Nearest_3_Spot', 'Nearest_3_Dist']] = df2.apply(
    lambda row: find_nearest_spots(row['Latitude'], row['Longitude']), axis=1)

优化方案 #

1. CPU端最优：基于BallTree的空间索引搜索 #

使用sklearn.neighbors.BallTree专门处理球面最近邻搜索，时间复杂度从O(m*n)降至O(m log n)，性能提升百倍以上。

import pandas as pd
import numpy as np
from sklearn.neighbors import BallTree

# 加载数据（替换为真实数据集）
df = pd.DataFrame(data)
df2 = pd.DataFrame(data_df2)

# 经纬度转弧度（haversine公式要求）
df_rad = np.radians(df[['Latitude', 'Longitude']].values)
df2_rad = np.radians(df2[['Latitude', 'Longitude']].values)

# 构建BallTree，使用haversine度量（球面距离）
tree = BallTree(df_rad, metric='haversine')

# 查询每个点的最近3个邻居，返回距离（弧度）和索引
distances_rad, indices = tree.query(df2_rad, k=3)

# 弧度转公里（地球平均半径6371km）
distances_km = distances_rad * 6371

# 整理结果到df2
for i in range(3):
    df2[f'Nearest_{i+1}_Spot'] = df.iloc[indices[:, i]]['City'].values
    df2[f'Nearest_{i+1}_Dist'] = distances_km[:, i]

2. CPU端次优：向量化Haversine计算 #

用numpy向量化实现Haversine距离公式，避免apply的循环开销，比原方案快10-20倍。

import pandas as pd
import numpy as np

def haversine_vectorized(lat1, lon1, lat2, lon2):
    # 转弧度
    lat1, lon1, lat2, lon2 = map(np.radians, [lat1, lon1, lat2, lon2])
    dlat = lat2 - lat1
    dlon = lon2 - lon1
    a = np.sin(dlat/2)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon/2)**2
    c = 2 * np.arcsin(np.sqrt(a))
    return c * 6371  # 转公里

# 提取df的经纬度数组
df_lat = df['Latitude'].values[:, np.newaxis]
df_lon = df['Longitude'].values[:, np.newaxis]

# 遍历df2的每个点，批量计算距离并取top3
results = []
for _, row in df2.iterrows():
    dists = haversine_vectorized(row['Latitude'], row['Longitude'], df_lat, df_lon).flatten()
    top3_idx = dists.argsort()[:3]
    top3_spots = df.iloc[top3_idx]['City'].values
    top3_dists = dists[top3_idx]
    results.append({
        'Nearest_1_Spot': top3_spots[0],
        'Nearest_1_Dist': top3_dists[0],
        'Nearest_2_Spot': top3_spots[1],
        'Nearest_2_Dist': top3_dists[1],
        'Nearest_3_Spot': top3_spots[2],
        'Nearest_3_Dist': top3_dists[2]
    })

df2 = pd.concat([df2, pd.DataFrame(results)], axis=1)

3. GPU加速：使用RAPIDS cuDF+cuML #

RAPIDS是NVIDIA免费GPU库，专为大规模数据处理优化，适合50万条级别的数据，性能比CPU方案再提升5-10倍。

import cudf
from cuml.neighbors import NearestNeighbors

# 用cuDF加载数据（支持从csv/parquet读取，或从pandas转换）
df_cudf = cudf.DataFrame(df)
df2_cudf = cudf.DataFrame(df2)

# 提取经纬度特征
X = df_cudf[['Latitude', 'Longitude']]
X_query = df2_cudf[['Latitude', 'Longitude']]

# 构建最近邻模型，使用haversine度量
knn = NearestNeighbors(n_neighbors=3, metric='haversine')
knn.fit(X)

# 查询距离和索引，注意haversine返回的距离是弧度，需转公里
distances_rad, indices = knn.kneighbors(X_query)
distances_km = distances_rad * 6371

# 映射城市名称并添加到结果
for i in range(3):
    df2_cudf[f'Nearest_{i+1}_Spot'] = df_cudf.iloc[indices[:, i]]['City'].values
    df2_cudf[f'Nearest_{i+1}_Dist'] = distances_km[:, i]

# 可选：转回pandas DataFrame
df2 = df2_cudf.to_pandas()

方案选择建议 #

• 无GPU时优先选BallTree方案，实现简单且性能最优；
• 有NVIDIA GPU时用RAPIDS方案，获得最大性能提升；
• 向量化方案适合无法安装sklearn的场景，性能介于原方案和BallTree之间。

内容的提问来源于stack exchange，提问作者sdave