使用 apply() 加速数据帧上的嵌套循环

Question

我有一个在 Python 中使用 Pandas 的数据框，其中每行包含纬度和经度坐标。我的目标是添加另一个名为“close_by”的列，其中包含数据集中 1 英里内的其他条目的数量，使用 hasrsine。

我看过其他类似问题的指南，例如：https://engineering.upside.com/a-beginners-guide-to-optimizing-pandas-code-for-speed-c09ef2c6a4d6 但它们涉及使用 df.apply() 更新每一行以添加坐标和一些静态每个定义点之间的距离。我没有任何运气找到或想出解决方案。

基本上，这就是我要优化的内容：

for index1, row1 in business_data.iterrows():
    for index2, row2 in business_data.iterrows():
        distance = mpu.haversine_distance((business_data.at[index1,'latitude'], business_data.at[index1,'longitude']), (business_data.at[index2,'latitude'], business_data.at[index2,'longitude']))
        distance = distance * 0.621371

        if distance <= 1:
            business_data.at[index1,'close_by'] = row1["close_by"] + 1

我有大约 50,000 行，在我的计算机上每行大约需要 5 秒。

感谢您的任何建议！

Answer 1

从外观上看，mpu.haversine_distance() 使用 math 而不是 numpy 函数，因此它不可矢量化。

改用this vectorized haversine distance function，您可以轻松地将问题向量化：

df = pd.DataFrame([
    {'latitude': 49.001, 'longitude': 11.0},
    {'latitude': 49.0, 'longitude': 11.0},
    {'latitude': 49.001, 'longitude': 11.001},
    {'latitude': -49.0, 'longitude': 11.0},
])


lon = df['longitude'].to_numpy()
lat = df['latitude'].to_numpy()

radius = 1.0

df['close_by'] = np.count_nonzero(haversine_np(lon, lat, lon[:, None], lat[:, None]) < radius, axis=0) - 1

df
#   latitude    longitude   nearby
# 0 49.001      11.000      2
# 1 49.000      11.000      2
# 2 49.001      11.001      2
# 3 -49.000     11.000      0