GTFS站点关联列表生成优化：DataFrame迭代耗时过长求助

问题描述 #

需要从GTFS数据库快照生成公交线路换乘关联列表，核心是识别可换乘的站点对（用于后续NetworkX图可视化），筛选规则为：

• 两个站点名称相同，视为可换乘
• 两个站点名称不同，但地理距离小于500米，视为可换乘

已预处理得到站点DataFrame：

import pandas as pd
stops = pd.DataFrame({'stop_id':  ['002133', '003002', '003118', '003209', '004521', '004716', '004903', '006390', '007177', '007289'], 
                      'stop_name':['Örs vezér tere M+H, déli tárolótér', 'Puskás Ferenc Stadion', 'Óbuda, Bogdáni út', 'Batthyány tér', 'Puskás Ferenc Stadion', 'ÉD metró járműtelep,porta', 'Örs vezér tere', 'Cinkota kocsiszín', 'Csepel kocsiszín', 'Vihar utca'],
                      'location': [(47.500366, 19.1357), (47.500368, 19.103406), (47.551471, 19.041971), (47.506776, 19.039318), (47.50017, 19.104773), (47.469651, 19.12909), (47.503585, 19.137192), (47.519345, 19.217072), (47.421498, 19.066247), (47.434399, 19.035664)]})

原方案采用嵌套迭代所有站点对，6000行数据时耗时极长——原因是O(n²)的暴力遍历加上每次调用geopy.distance.geodesic的高开销，导致计算量呈指数级增长。

优化方案 #

核心思路 #

拆分需求为两个独立部分分别处理，用空间索引替代暴力距离计算，将时间复杂度从O(n²)降至O(n log n)：

1. 同名称站点对：通过分组+组合生成，避免无效遍历
2. 异名称近距站点对：用KDTree空间索引快速定位邻域站点，减少距离计算次数

具体实现代码 #

1. 导入依赖库

import pandas as pd
import numpy as np
from scipy.spatial import KDTree
from itertools import combinations

2. 预处理坐标数据

将经纬度转换为弧度（适配Haversine距离计算）：

# 拆分经纬度并转换为弧度
stops['lat_rad'] = np.radians(stops['location'].apply(lambda x: x[0]))
stops['lon_rad'] = np.radians(stops['location'].apply(lambda x: x[1]))
coords = stops[['lat_rad', 'lon_rad']].values

# 地球半径（单位：米）
EARTH_RADIUS = 6371000

3. 生成同名称站点对

按站点名称分组，生成组内所有两两组合：

same_name_pairs = []
for name, group in stops.groupby('stop_name'):
    if len(group) >= 2:
        # 生成组内无序站点对，避免重复（如A-B和B-A只保留一次）
        for pair in combinations(group.itertuples(), 2):
            same_name_pairs.append({
                'stop1': pair[0].stop_name,
                'stop2': pair[1].stop_name,
                'stop1_id': pair[0].stop_id,
                'stop2_id': pair[1].stop_id,
                'stop1_location': pair[0].location,
                'stop2_location': pair[1].location
            })

4. 生成异名称近距站点对

用KDTree快速查找500米内的邻域站点：

# 构建KDTree，使用Haversine距离计算
kdtree = KDTree(coords, metric='haversine')

# 将500米转换为弧度距离（距离/地球半径）
radius = 500 / EARTH_RADIUS
# 查找每个站点的邻域索引
nearby_indices = kdtree.query_ball_point(coords, radius)

near_diff_name_pairs = []
for i, indices in enumerate(nearby_indices):
    for j in indices:
        # 只保留i<j的对，避免重复
        if i < j and stops.iloc[i]['stop_name'] != stops.iloc[j]['stop_name']:
            row1 = stops.iloc[i]
            row2 = stops.iloc[j]
            near_diff_name_pairs.append({
                'stop1': row1['stop_name'],
                'stop2': row2['stop_name'],
                'stop1_id': row1['stop_id'],
                'stop2_id': row2['stop_id'],
                'stop1_location': row1['location'],
                'stop2_location': row2['location']
            })

5. 合并结果

# 合并两类站点对
all_transfer_pairs = same_name_pairs + near_diff_name_pairs

# 转换为DataFrame方便后续处理
pairs_df = pd.DataFrame(all_transfer_pairs)
print(pairs_df)

NetworkX图可视化适配 #

直接将站点作为节点、换乘对作为边导入图：

import networkx as nx
import matplotlib.pyplot as plt

# 创建无向图
G = nx.Graph()

# 添加所有站点节点
for _, row in stops.iterrows():
    G.add_node(
        row['stop_id'],
        name=row['stop_name'],
        location=row['location']
    )

# 添加换乘边
for pair in all_transfer_pairs:
    G.add_edge(
        pair['stop1_id'],
        pair['stop2_id'],
        edge_type='transfer'
    )

# 简单可视化
nx.draw(G, with_labels=True, node_size=1000, font_size=8)
plt.show()

内容的提问来源于stack exchange，提问作者Erik Gebhard