如何用Python正则表达式为Networkx图设置主/次边颜色（橙/黄）

步骤说明 #

1. 解析边数据：将提供的JSON格式边数据转换为Python可处理的字典列表。
2. 正则匹配边类型：通过正则表达式从每条边的label字段中提取primary或secondary标识。
3. 面向对象封装图逻辑：创建可视化类，负责加载边数据、设置边颜色属性，以及执行可视化渲染。
4. 渲染彩色边：利用NetworkX的绘图接口，根据边的颜色属性设置对应颜色。

完整代码实现 #

import re
import json
import networkx as nx
import matplotlib.pyplot as plt

class GraphVisualizer:
    def __init__(self):
        self.graph = nx.Graph()
        # 定义边类型与颜色的映射
        self.color_map = {
            "primary": "orange",
            "secondary": "yellow"
        }
        # 匹配label中边类型的正则表达式（忽略大小写）
        self.type_pattern = re.compile(r'(primary|secondary)', re.IGNORECASE)

    def load_edges(self, edge_data_str):
        # 将多行JSON字符串转为字典列表
        edges = [json.loads(line.strip()) for line in edge_data_str.split('\n') if line.strip()]
        
        for edge in edges:
            nodes = edge["nodes_in_edge"]
            # 用正则匹配边类型
            match = self.type_pattern.search(edge["label"])
            edge_type = match.group(1).lower() if match else "default"
            # 获取对应颜色，默认灰色
            color = self.color_map.get(edge_type, "gray")
            
            # 处理多节点边：拆分为连续边段
            if len(nodes) >= 2:
                for i in range(len(nodes)-1):
                    self.graph.add_edge(
                        nodes[i], 
                        nodes[i+1], 
                        color=color, 
                        label=edge["label"], 
                        width=edge["width"]
                    )

    def visualize(self):
        # 提取边的颜色和宽度参数
        edge_colors = [self.graph[u][v]['color'] for u, v in self.graph.edges()]
        edge_widths = [float(self.graph[u][v]['width'].replace("'", "")) for u, v in self.graph.edges()]
        
        # 绘制图
        pos = nx.spring_layout(self.graph)
        nx.draw(
            self.graph, 
            pos, 
            edge_color=edge_colors, 
            width=edge_widths, 
            with_labels=True, 
            node_color='lightblue'
        )
        # 添加边标签
        edge_labels = nx.get_edge_attributes(self.graph, 'label')
        nx.draw_networkx_edge_labels(self.graph, pos, edge_labels=edge_labels)
        
        plt.show()

# 示例使用
if __name__ == "__main__":
    # 你的边数据
    edge_data = '''{"id": 764046068, "label": "edge 6 secondary", "nodes_in_edge": [5653832464, 7138043089, 7138030356], "width": "2'", "speed": "20 mph"}
{"id": 1053591753, "label": "edge 8 primary", "nodes_in_edge": [5475082077, 9682439803], "width": "4'", "speed": "50 mph"}'''
    
    visualizer = GraphVisualizer()
    visualizer.load_edges(edge_data)
    visualizer.visualize()

关键细节说明 #

• 正则匹配：re.compile(r'(primary|secondary)', re.IGNORECASE) 可以兼容大小写形式的标识，避免匹配失败。
• 多节点边处理：将包含3个及以上节点的边拆分为连续边段，确保NetworkX能正常渲染。
• 颜色映射：通过字典color_map统一管理类型与颜色的对应关系，后续修改或扩展更便捷。

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