基于Python Pandas DataFrame实现ID关联遍历生成新数据表

针对你的需求，SQL循环在数据量大时效率低下，推荐用图论算法来处理，这是解决这类关联遍历问题的高效方案，下面是具体实现步骤：

1. 工具选择 #

用networkx库处理图的连通性，它的算法时间复杂度接近线性，远快于逐行循环。如果没安装，先执行：

pip install networkx

2. 具体实现代码 #

假设你的原始DataFrame名为df，结构示例如下：

情况1：获取所有双向关联ID（比如A和B、C互相视为关联） #

import pandas as pd
import networkx as nx

# 1. 构建无向图
G = nx.Graph()
# 添加所有节点（New ID）
G.add_nodes_from(df['New ID'].unique())
# 添加边：New ID 与对应的 Prev ID（排除空值）
valid_edges = df.dropna(subset=['Prev ID'])[['New ID', 'Prev ID']].values.tolist()
G.add_edges_from(valid_edges)

# 2. 生成ID-YEAR映射字典，快速查找年份
id_year_map = df.set_index('New ID')['YEAR'].to_dict()

# 3. 遍历每个连通分量，生成结果数据
result_data = []
for component in nx.connected_components(G):
    component_ids = list(component)
    # 每个Selected ID对应分量内所有Related ID
    for selected_id in component_ids:
        for related_id in component_ids:
            result_data.append({
                'Selected ID': selected_id,
                'Related ID': related_id,
                'YEAR': id_year_map[related_id]
            })

# 转换为最终DataFrame
result_df = pd.DataFrame(result_data)

情况2：单向追溯关联ID（从Selected ID出发，仅通过Prev ID向上遍历，比如A的关联ID是A、B、C） #

如果是单向的前驱追溯，需要构建有向图：

import pandas as pd
import networkx as nx

# 1. 构建有向图（边方向：New ID → Prev ID，代表从当前ID指向它的前驱）
G = nx.DiGraph()
G.add_nodes_from(df['New ID'].unique())
valid_edges = df.dropna(subset=['Prev ID'])[['New ID', 'Prev ID']].values.tolist()
G.add_edges_from(valid_edges)

# 2. 生成ID-YEAR映射字典
id_year_map = df.set_index('New ID')['YEAR'].to_dict()

# 3. 遍历每个Selected ID，查找所有可达的关联ID（包括自己）
result_data = []
for selected_id in df['New ID'].unique():
    # 获取从当前ID出发能到达的所有节点（包括自身）
    related_ids = nx.single_source_shortest_path_nodes(G, selected_id)
    for related_id in related_ids:
        result_data.append({
            'Selected ID': selected_id,
            'Related ID': related_id,
            'YEAR': id_year_map.get(related_id, None)  # 兼容可能不存在的ID
        })

result_df = pd.DataFrame(result_data)

为什么这个方法更快？ #

SQL循环是逐行迭代匹配，时间复杂度约为O(n²)，数据量越大越慢；而networkx的连通性算法基于图论，时间复杂度为O(n+m)（n是节点数，m是边数），属于线性级别，处理百万级数据也能保持高效。

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