如何基于连续端点对Pandas DataFrame中的线段进行排序

问题描述 #

给定一个Pandas DataFrame，每行包含row_id、线段左边界坐标left_coord和右边界坐标right_coord。这些线段可以按首尾相连的顺序排列（即第i条线段的right_coord与第i+1条的left_coord完全重合）。需要实现通用的排序逻辑，不能依赖单纯按left_coord排序的方法（该方法仅在特定场景有效，不具备通用性）。

示例输入：

import pandas as pd
df = pd.DataFrame({
    'row_id':[1,2,3,4], 
    'left_coord': [(0,0), (0,1), (4,4), (1,1)], 
    'right_coord':[(0,1),(1,1),(5,7),(4,4)]
})

初始DataFrame：

row_id left_coord right_coord
0       1     (0, 0)      (0, 1)
1       2     (0, 1)      (1, 1)
2       3     (4, 4)      (5, 7)
3       4     (1, 1)      (4, 4)

期望输出：

row_id left_coord right_coord
0       1     (0, 0)      (0, 1)
1       2     (0, 1)      (1, 1)
3       4     (1, 1)      (4, 4)
2       3     (4, 4)      (5, 7)

通用解决方案 #

核心思路是先找到线段链的起点，再通过坐标映射依次匹配后续线段，最终生成有序的DataFrame。具体步骤如下：

• 构建left_coord到对应行数据的映射表，实现快速查找
• 定位链的起点：找到left_coord未出现在任何线段right_coord集合中的线段（无前置线段的起始点）
• 从起点开始，按right_coord匹配下一条线段，直到所有线段都被收集
• 将有序线段列表转换为DataFrame

代码实现 #

import pandas as pd

def sort_connected_segments(df):
    # 构建left_coord到行数据的映射，方便快速查找
    left_to_row = {row['left_coord']: row for _, row in df.iterrows()}
    
    # 获取所有right_coord的集合，用于找起点
    right_coords = set(df['right_coord'])
    
    # 筛选起始线段：left_coord不在right_coords中的行
    start_row = df[~df['left_coord'].isin(right_coords)].iloc[0]
    
    # 初始化有序列表
    ordered_rows = [start_row]
    
    # 迭代匹配后续线段
    while len(ordered_rows) < len(df):
        current_right = ordered_rows[-1]['right_coord']
        next_row = left_to_row[current_right]
        ordered_rows.append(next_row)
    
    # 转换为DataFrame并重置索引
    return pd.DataFrame(ordered_rows).reset_index(drop=True)

# 测试示例
df = pd.DataFrame({
    'row_id':[1,2,3,4], 
    'left_coord': [(0,0), (0,1), (4,4), (1,1)], 
    'right_coord':[(0,1),(1,1),(5,7),(4,4)]
})

sorted_df = sort_connected_segments(df)
print(sorted_df)

代码说明 #

• 映射表构建：通过字典直接将left_coord映射到行数据，后续查找下一条线段的时间复杂度为O(1)
• 起点定位：利用集合的快速查找特性，高效筛选出无前置线段的起始点
• 迭代排序：从起点出发，沿着right_coord -> left_coord的关联依次匹配，确保线段首尾相连
• 无冗余处理：题目保证线段可形成完整链，因此无需额外处理循环情况（若实际场景有循环风险，可添加已处理坐标集合进行校验）

验证结果 #

运行代码后输出与期望完全一致：

row_id left_coord right_coord
0       1     (0, 0)      (0, 1)
1       2     (0, 1)      (1, 1)
2       4     (1, 1)      (4, 4)
3       3     (4, 4)      (5, 7)

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