绘制层次结构图

news/2024/11/15 22:42:53/文章来源:https://www.cnblogs.com/epicblue/p/18548794

绘制层次结构图

WPS的智能图形要收费，先做个免费的不美观的版本。

基于
matplotlib，networkx，graphviz，pydot

按需修改
输入内容
input_data 为输入的文本。

外观
rankdir 为指定方向。
mpatches.Rectangle 为节点外形。

比例
缩放matplotlib窗口，调整节点长宽。
调整字体大小，当前为 plt.text(fontsize=10）。

代码
import matplotlib.pyplot as plt
import networkx as nx
import matplotlib.patches as mpatches

plt.rcParams['font.sans-serif'] = ['SimHei'] # 正常显示中文
plt.rcParams['axes.unicode_minus'] = False # 正常显示负号

def parse_hierarchy(input_str):
"""
解析组织架构的层级输入，生成树结构。
"""
lines = input_str.strip().split("\n")
root = None
stack = [] # 用于追踪当前的父层级节点
hierarchy = {}

for line in lines:# 计算当前行的缩进级别stripped_line = line.lstrip("- ").strip()level = (len(line) - len(line.lstrip("- "))) // 2# 创建当前节点if root is None:root = stripped_linehierarchy[root] = []stack.append((root, level))else:while stack and stack[-1][1] >= level:  # 回退到上一层级节点stack.pop()if stack:parent, _ = stack[-1]hierarchy[parent].append(stripped_line)else:# 如果栈为空但仍有节点，则说明输入格式有问题raise ValueError(f"错误的层级结构：无法找到父节点来连接 {stripped_line}")hierarchy[stripped_line] = []stack.append((stripped_line, level))return root, hierarchy

def plot_organization_chart(root, hierarchy, rankdir = "TB"):
G = nx.DiGraph()

def add_edges(parent, children):for child in children:G.add_edge(parent, child)add_edges(child, hierarchy.get(child, []))add_edges(root, hierarchy[root])# 创建一个 Pydot 的图对象
dot = nx.drawing.nx_pydot.to_pydot(G)# 设置图的方向
dot.set_rankdir(rankdir)  # 'LR' 为从左到右，'TB' 为从上到下# Pydot 的图对象 倒腾到 G
G = nx.drawing.nx_pydot.from_pydot(dot)# 使用层次布局定位节点
pos = nx.drawing.nx_pydot.graphviz_layout(G, prog='dot')# 缩小节点间距
def scale_pos(pos, scale_x=1.0, scale_y=1.0):return {node: (x * scale_x, y * scale_y) for node, (x, y) in pos.items()}# 调整缩放比例以减少节点之间的间距
pos = scale_pos(pos, scale_x=0.5, scale_y=0.5)plt.figure(figsize=(10, 6))# 创建长方形节点
def draw_rect_node(node, pos):x, y = pos[node]rect_width, rect_height = 40, 15  # 矩形宽度和高度rect = mpatches.Rectangle((x - rect_width / 2, y - rect_height / 2),rect_width, rect_height,edgecolor='black', facecolor='lightblue', alpha=0.8)plt.gca().add_patch(rect)  # 添加矩形到当前轴plt.text(x, y, node, ha='center', va='center', fontsize=10, fontweight='bold', color='black')# 绘制节点和矩形
rect_width, rect_height = 40, 15
for node in G.nodes():draw_rect_node(node, pos)for parent, child in G.edges():x0, y0 = pos[parent]x1, y1 = pos[child]if rankdir == "TB" or rankdir == "BT":# 计算起点和终点，避开矩形区域if y0 > y1:  # 从上到下start_y = y0 - rect_height / 2end_y = y1 + rect_height / 2else:  # 从下到上start_y = y0 + rect_height / 2end_y = y1 - rect_height / 2# 保持 x 坐标不变start_x = x0end_x = x1y_mid = (start_y + end_y) / 2  # 中间的水平线 y 坐标# 绘制边plt.plot([start_x, start_x], [start_y, y_mid], "k-", linewidth=0.8)  # 垂直线plt.plot([start_x, end_x], [y_mid, y_mid], "k-", linewidth=0.8)      # 水平线plt.plot([end_x, end_x], [y_mid, end_y], "k-", linewidth=0.8)        # 垂直线else:# 计算起点和终点，避开矩形区域if x0 < x1:  # 从左到右start_x = x0 + rect_width / 2end_x = x1 - rect_width / 2else:  # 从右到左 (虽然这里不太可能出现，但为了代码的完整性，还是加上)start_x = x0 - rect_width / 2end_x = x1 + rect_width / 2# 保持 y 坐标不变start_y = y0end_y = y1x_mid = (start_x + end_x) / 2  # 中间的垂直线 x 坐标# 绘制边plt.plot([start_x, x_mid], [start_y, start_y], "k-", linewidth=0.8)  # 水平线plt.plot([x_mid, x_mid], [start_y, end_y], "k-", linewidth=0.8)      # 垂直线plt.plot([x_mid, end_x], [end_y, end_y], "k-", linewidth=0.8)        # 水平线plt.title("Organization Chart", fontsize=14)
plt.axis("off")
plt.tight_layout()
plt.show()