如何用NetworkX递归查找DataFrame中指定方法的连通组件

问题背景 #

给定如下DataFrame，main列代表方法，component列是该方法包含的组件（组件也可作为其他方法，即子方法）：

import numpy as np
import pandas as pd

df = pd.DataFrame(
    {
        "main": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
        "component": [
            [1, 2],
            [np.nan],
            [3, 8],
            [np.nan],
            [1, 5, 6],
            [np.nan],
            [7],
            [np.nan],
            [9, 10],
            [np.nan],
            [np.nan],
        ],
    }
)

需求：输入指定方法，找出其所有连通的子方法/组件。例如指定方法0时，期望输出：

target = pd.DataFrame({
    "main": [0, 0, 2, 2, 8, 8],
    "component": [1, 2, 3, 8, 9, 10]
})

同时需要生成对应的可视化图。

解决方案 #

借助networkx构建方法与组件的关联图，通过连通子图功能提取目标方法的所有关联节点，再整理成目标格式，最后用graphviz可视化。

1. 数据预处理 #

先展开component列并移除无组件的条目，同时将组件类型转为整数：

df_exploded = df.explode(column="component").dropna(subset="component")
df_exploded['component'] = df_exploded['component'].astype(int)

2. 构建关联图 #

用networkx创建无向图，添加所有方法与组件的关联边：

import networkx as nx

G = nx.Graph()
# 提取所有(main, component)配对作为边
edges = list(df_exploded[['main', 'component']].itertuples(index=False, name=None))
G.add_edges_from(edges)

3. 获取目标方法的连通组件 #

编写函数，输入目标方法，返回其所有连通组件的DataFrame：

def get_connected_components(target_main):
    # 获取目标节点所在的连通子图节点集合
    connected_nodes = nx.node_connected_component(G, target_main)
    # 筛选出属于该连通子图的所有边
    connected_edges = df_exploded[
        df_exploded['main'].isin(connected_nodes) & 
        df_exploded['component'].isin(connected_nodes)
    ]
    return connected_edges.reset_index(drop=True)

# 测试：获取方法0的连通组件
target_df = get_connected_components(0)
print(target_df)

执行后输出的target_df与需求中的target完全一致。

4. 可视化连通关系 #

用graphviz生成可视化图，展示目标方法的所有连通层级：

import graphviz

def visualize_connected_graph(target_main):
    connected_nodes = nx.node_connected_component(G, target_main)
    subgraph_edges = df_exploded[
        df_exploded['main'].isin(connected_nodes) & 
        df_exploded['component'].isin(connected_nodes)
    ]
    
    # 配置图属性：横向布局、节点间距
    graph_attr = dict(rankdir="LR", nodesep="0.2")
    g = graphviz.Digraph(graph_attr=graph_attr)
    
    # 添加节点
    for node in connected_nodes:
        g.node(str(node), shape="box", style="filled", height="0.35")
    
    # 添加边
    for _, row in subgraph_edges.iterrows():
        g.edge(str(row['main']), str(row['component']))
    
    return g

# 生成方法0的可视化图（Jupyter环境可直接显示，也可保存为文件）
viz_graph = visualize_connected_graph(0)
viz_graph
# 保存为PDF：viz_graph.render("method_0_connections")

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