基于Python实现指定点100000米范围内ID筛选的代码问题排查

现有如下结构的DataFrame：

需要编写函数，为每条记录计算其100000米半径范围内的记录数量及对应ID。以下是编写的代码，但未得到预期结果：

import numpy as np
from scipy.spatial import cKDTree
from math import radians, sin, cos, sqrt, atan2

R = 6371000  # 地球半径（米）

def haversine(lat1, lon1, lat2, lon2):
    phi1, phi2 = radians(lat1), radians(lat2)
    delta_phi = radians(lat2 - lat1)
    delta_lambda = radians(lon2 - lon1)
    
    a = sin(delta_phi / 2)**2 + cos(phi1) * cos(phi2) * sin(delta_lambda / 2)**2
    c = 2 * atan2(sqrt(a), sqrt(1 - a))
    
    return R * c

def lat_lon_to_cartesian(lat, lon):
    lat, lon = np.radians(lat), np.radians(lon)
    x = R * np.cos(lat) * np.cos(lon)
    y = R * np.cos(lat) * np.sin(lon)
    z = R * np.sin(lat)
    return x, y, z

# 为DataFrame添加笛卡尔坐标
df['x'], df['y'], df['z'] = zip(*df.apply(lambda row: lat_lon_to_cartesian(row['lat'], row['lon']), axis=1))

# 构建KD-Tree用于快速空间索引
tree = cKDTree(df[['x', 'y', 'z']])

# 使用KD-Tree查找100000米半径内的点
def points_within_radius_kdtree(lat, lon, radius=100000):
    x, y, z = lat_lon_to_cartesian(lat, lon)
    indices = tree.query_ball_point([x, y, z], radius / R)  # 调整半径适配笛卡尔空间
    nearby_points = df.iloc[indices].copy()
    nearby_points['distance'] = nearby_points.apply(lambda row: haversine(lat, lon, row['lat'], row['lon']), axis=1)
    return nearby_points[nearby_points['distance'] <= radius]

# 为每条记录计算100000米半径内的记录数量及ID
def calculate_nearby_counts(df, radius=100000):
    counts = []
    ids = []
    
    for idx, row in df.iterrows():
        nearby_points = points_within_radius_kdtree(row['lat'], row['lon'], radius)
        count = len(nearby_points) - 1  # 减去自身
        nearby_ids = nearby_points[nearby_points['ID'] != row['ID']]['ID'].tolist()  # 排除自身
        counts.append(count)
        ids.append(nearby_ids)
    
    df['nearby_count'] = counts
    df['nearby_ids'] = ids
    return df

result_df = calculate_nearby_counts(df)
print(result_df)

预期输出示例（以ID=0的记录为例）：

代码存在3个核心问题：

1. 列名不匹配：DataFrame中纬度列名为latitude，但代码中多次错误使用row['lat']，触发KeyError。
2. KD-Tree半径参数错误：转换后的笛卡尔坐标以米为单位，无需将半径除以地球半径，直接传入100000即可。
3. 逻辑冗余：计算附近点数量时，先减1再过滤自身，不如直接过滤后再统计更清晰。

修正后的代码如下：

import numpy as np
import pandas as pd
from scipy.spatial import cKDTree
from math import radians, sin, cos, sqrt, atan2

R = 6371000  # 地球半径（米）

def haversine(lat1, lon1, lat2, lon2):
    phi1, phi2 = radians(lat1), radians(lat2)
    delta_phi = radians(lat2 - lat1)
    delta_lambda = radians(lon2 - lon1)
    
    a = sin(delta_phi / 2)**2 + cos(phi1) * cos(phi2) * sin(delta_lambda / 2)**2
    c = 2 * atan2(sqrt(a), sqrt(1 - a))
    
    return R * c

def lat_lon_to_cartesian(lat, lon):
    lat, lon = np.radians(lat), np.radians(lon)
    x = R * np.cos(lat) * np.cos(lon)
    y = R * np.cos(lat) * np.sin(lon)
    z = R * np.sin(lat)
    return x, y, z

# 初始化示例DataFrame
df = pd.DataFrame({
    'latitude': [-22.582779, -22.582575, 41.758910, 17.758527, -22.582699],
    'longitude': [29.080456, 29.080794, -53.626698, -2.443443, 29.080455],
    'ID': [0, 1, 2, 3, 4]
})

# 修正列名，添加笛卡尔坐标
df['x'], df['y'], df['z'] = zip(*df.apply(lambda row: lat_lon_to_cartesian(row['latitude'], row['longitude']), axis=1))

# 构建KD-Tree
tree = cKDTree(df[['x', 'y', 'z']])

def points_within_radius_kdtree(lat, lon, radius=100000):
    x, y, z = lat_lon_to_cartesian(lat, lon)
    # 直接使用米为单位的半径
    indices = tree.query_ball_point([x, y, z], radius)
    nearby_points = df.iloc[indices].copy()
    # 计算实际球面距离，过滤误差点
    nearby_points['distance'] = nearby_points.apply(lambda row: haversine(lat, lon, row['latitude'], row['longitude']), axis=1)
    return nearby_points[nearby_points['distance'] <= radius]

def calculate_nearby_counts(df, radius=100000):
    counts = []
    ids = []
    
    for idx, row in df.iterrows():
        nearby_points = points_within_radius_kdtree(row['latitude'], row['longitude'], radius)
        # 直接过滤自身记录
        filtered_points = nearby_points[nearby_points['ID'] != row['ID']]
        counts.append(len(filtered_points))
        ids.append(filtered_points['ID'].tolist())
    
    # 添加预期列名
    df['#counts within 100000m'] = counts
    df['ids within 100000m'] = ids
    # 保留需要的结果列
    return df[['latitude', 'longitude', 'ID', '#counts within 100000m', 'ids within 100000m']]

result_df = calculate_nearby_counts(df)
print(result_df)

修正后的代码输出符合预期：

latitude  longitude  ID  #counts within 100000m ids within 100000m
0 -22.582779   29.080456   0                        2             [1, 4]
1 -22.582575   29.080794   1                        2             [0, 4]
2  41.758910  -53.626698   2                        0                []
3  17.758527   -2.443443   3                        0                []
4 -22.582699   29.080455   4                        2             [0, 1]

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