繪製層次結構圖

lusonixs發表於2024-11-15

繪製層次結構圖

Word的SmartArt挺好，先來個免費的不美觀的版本。

基於
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_line
            hierarchy[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 / 2
                end_y = y1 + rect_height / 2
            else:  # 從下到上
                start_y = y0 + rect_height / 2
                end_y = y1 - rect_height / 2

            # 保持 x 座標不變
            start_x = x0
            end_x = x1

            y_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 / 2
                end_x = x1 - rect_width / 2
            else:  # 從右到左 (雖然這裡不太可能出現，但為了程式碼的完整性，還是加上)
                start_x = x0 - rect_width / 2
                end_x = x1 + rect_width / 2

            # 保持 y 座標不變
            start_y = y0
            end_y = y1

            x_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()


# 輸入的層次結構文字
input_data = """
頂層節點
- 一級節點1
- - 二級節點1
- - 二級節點2
- - 二級節點3
- - 二級節點4
- - - 三級節點1
- - - 三級節點2
- - - 三級節點3
- 一級節點2
- - 二級節點5
- - 二級節點6
"""

try:
    root, hierarchy = parse_hierarchy(input_data)
    plot_organization_chart(root, hierarchy, "LR")
except ValueError as e:
    print(f"輸入解析錯誤：{e}")

相關文章