python - 计算 Pandas 地理密度的有效方法？

Question

我有大量与美国快餐店对应的经度和纬度数据列表。对于每个快餐店，我想知道 5 英里以内还有多少其他快餐店。我可以像这样使用 Geopy 在 Pandas 中计算这个(DataFrame 中的每一行都是一个不同的快餐店):

import pandas as pd
import geopy.distance

df = pd.DataFrame({'Fast Food Place':[1,2,3], 'Lat':[33,34,35], 'Lon':[42,43,44]})

for index1, row1 in df.iterrows():
    num_fastfood = 0

    for index2, row2 in df.iterrows():
        # calculate distance in miles between longitude and latitude
        dist = geopy.distance.VincentyDistance(row1[['Lat','Lon']],
                                               row2[['Lat','Lon']]).miles

        # if fast food is within 5 miles, increment num_fastfood
        if dist < 5: # if less than five miles
            num_fastfood = num_fastfood + 1

    df.loc[index1, 'num_fastfood_5miles'] = num_fastfood - 1 # (subtract 1 to exclude self)

但这在非常大的数据集(即 50,000 行)上非常慢。我考虑过使用 KDTree 进行搜索，但好奇其他人是否有更快的方法？

Answer 1

scipy.spatial.cKDTree 的实现:

from scipy.spatial import cKDTree

def find_neighbours_within_radius(xy, radius):
    tree = cKDTree(xy)
    within_radius = tree.query_ball_tree(tree, r=radius)
    return within_radius

def flatten_nested_list(nested_list):
    return [item for sublist in nested_list for item in sublist]

def total_neighbours_within_radius(xy, radius):
    neighbours = find_neighbours_within_radius(xy, radius)
    return len(flatten_nested_list(neighbours))