基于自定义函数的行对比：大DataFrame高效匹配相似行程位置

针对你的大型DataFrame，逐行循环对比的效率极低，下面是基于空间索引的高效实现方案，适合大数据量场景：

核心思路 #

用**空间索引树（KDTree/BallTree）**替代全量逐行对比，把时间复杂度从O(n²)降到O(n log n)，同时用向量运算替代Python循环，大幅提升速度。

步骤1：统一坐标距离计算逻辑 #

你的check函数本质是判断两点是否在指定半径内，分两种场景处理：

• 如果是平面坐标（比如投影后的笛卡尔坐标）：用欧氏距离
• 如果是经纬度（地理坐标）：用Haversine公式计算球面距离（更准确）

步骤2：完整代码实现 #

假设你的DataFrame名为df，先导入依赖：

import pandas as pd
import numpy as np
from scipy.spatial import KDTree
# 地理坐标专用
from sklearn.neighbors import BallTree

预处理数据 #

给每个行程生成唯一标识，方便后续关联结果：

# 生成唯一行程ID：车辆ID_行程ID
df['trip_id'] = df['Vehicle'].astype(str) + '_' + df['trip'].astype(str)

场景1：平面坐标（KDTree） #

# 提取坐标数组
coords = df[['longitude', 'latitude']].values

# 构建KDTree索引
tree = KDTree(coords)

# 设置半径（单位需和坐标一致，比如米）
radius = 100.0  # 示例：100米

# 查询每个点的邻域，排除自身
indices = tree.query_ball_point(coords, r=radius)
similar_indices = [idx[idx != i] for i, idx in enumerate(indices)]

场景2：经纬度地理坐标（BallTree） #

地理坐标需转成弧度，用Haversine距离计算：

# 经纬度转弧度（BallTree要求纬度在前）
coords_rad = np.radians(df[['latitude', 'longitude']].values)

# 构建BallTree，指定haversine度量
tree = BallTree(coords_rad, metric='haversine')

# 半径转弧度：地球平均半径6371千米，这里假设半径是1千米
radius_km = 1.0
radius_rad = radius_km / 6371.0

# 查询邻域，排除自身
indices = tree.query_radius(coords_rad, r=radius_rad)
similar_indices = [idx[idx != i] for i, idx in enumerate(indices)]

映射结果到DataFrame #

把查询到的索引转成对应的车辆/行程信息：

# 每个行程对应的相似行程ID列表
df['similar_trips'] = [df.iloc[idx]['trip_id'].tolist() for idx in similar_indices]

# 每个行程对应的相似车辆ID列表（去重）
df['similar_vehicles'] = [df.iloc[idx]['Vehicle'].unique().tolist() for idx in similar_indices]

步骤3：适配自定义check函数 #

如果你的check函数有特殊逻辑（比如额外条件），可以改成批量向量运算的形式，示例：

def batch_custom_check(main_coords, compare_coords, radius, main_vehicles, compare_vehicles):
    # 计算欧氏距离
    distances = np.linalg.norm(main_coords[:, np.newaxis] - compare_coords, axis=2)
    # 自定义条件：距离小于半径 + 不是同一车辆
    mask = (distances <= radius) & (main_vehicles[:, np.newaxis] != compare_vehicles)
    # 返回每个主点符合条件的索引
    return [np.where(row)[0] for row in mask]

# 调用示例
similar_indices = batch_custom_check(
    coords, coords, radius,
    df['Vehicle'].values, df['Vehicle'].values
)

关键优势 #

• 空间索引树避免了全量对比，大数据量下速度提升10~100倍
• 向量运算完全利用NumPy的C底层实现，比Python循环快几个数量级

内容的提问来源于stack exchange，提问作者Mohammad.sh