使用KDtree查找最近networkx最近节点

使用KDtree查找最近networkx最近节点

from sklearn.neighbors import KDTree
import networkx as nx
from geopy.distance import geodesic

def create_kdtree(G):
    coordinates = []
    node_list = []
    for node, data in G.nodes(data=True):
        coordinates.append([data['lat'], data['lon']])
        node_list.append(node)
        
    kd_tree = KDTree(coordinates)
    
    return kd_tree, node_list

def find_nearest_node(G, kd_tree, node_list, point):
    dist, ind = kd_tree.query([point], k=1)
    nearest_node = node_list[ind[0][0]]
    
    # Get coordinates of nearest node
    nearest_point = (G.nodes[nearest_node]['lat'], G.nodes[nearest_node]['lon'])
    
    # Calculate geodesic distance using geopy, 保留两位小数
    distance = round(geodesic(point, nearest_point).meters,2)
    
    return nearest_node, distance

def find_nearest_nodes(G, kd_tree, node_list, gdf):
    # 批次查找最近的节点
    nearest_nodes = []
    distances = []
    for point in gdf.geometry:
        nearest_node, distance = find_nearest_node(G, kd_tree, node_list, (point.y, point.x))
        nearest_nodes.append(nearest_node)
        distances.append(distance)
    gdf['nearest_node'] = nearest_nodes
    gdf['distance'] = distances
    return gdf