基于含FROM/TO列的无序DataFrame生成正确连接路径

给定一个包含from和to列的无序pandas DataFrame：

import pandas as pd

df = pd.DataFrame({
    'from': ['Start', '21', '73', 'Start', '55', '1', '2', '3'],
    'to': ['21', '73', '55', '1', '54', '2', '3', '4']
})

对应的表格如下：

需求 #

生成完整的连接路径列表：以Start为起点，依次匹配from列的当前节点与to列的下一个节点，直到无法匹配为止。已知连接为无环无分支结构，期望输出为：

[['Start', '21', '73', '55', '54'], ['Start', '1', '2', '3', '4']]

解决方案 #

# 构建from到to的映射字典，快速查找下一个节点
path_map = df.set_index('from')['to'].to_dict()

# 获取所有从Start出发的第一个节点
start_nodes = df[df['from'] == 'Start']['to'].tolist()

result_paths = []
for node in start_nodes:
    current_path = ['Start', node]
    next_node = path_map.get(node)
    # 顺着映射遍历直到没有下一个节点
    while next_node is not None:
        current_path.append(next_node)
        next_node = path_map.get(next_node)
    result_paths.append(current_path)

print(result_paths)

说明 #

1. 先将DataFrame转换为from到to的字典映射，实现O(1)时间复杂度的节点查找
2. 提取所有从Start出发的初始节点，每个节点对应一条独立路径
3. 对每条路径，持续查找并添加后续节点，直到找不到下一个节点为止，最终收集所有完整路径

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