需求：基于matplotlib开发二叉搜索树的自定义缺角Sunburst图表

要实现类似Plotly旭日图但带有缺角的效果，我们可以通过Matplotlib的patches.Wedge手动绘制多层环形区域，结合二叉搜索树的层级结构分配角度。以下是具体实现步骤和代码：

核心思路 #

1. 将二叉搜索树转换为层级化数据，记录每个节点的父节点、所在层级、权重（用于分配角度占比）
2. 预留缺角角度（比如60°），将剩余的有效区域作为绘制范围，所有节点的角度按权重比例分配到该区域
3. 按层级绘制环形楔形，每层对应树的一个深度，不同层级的环形宽度不同
4. 为每个节点添加标签，增强可读性

代码实现 #

1. 定义二叉搜索树结构 #

先构造一个示例家谱树，每个节点包含id、parent_id、name、weight（用于角度分配，比如代表家庭成员数量）：

tree_data = [
    {"id": 0, "parent_id": None, "name": "祖先", "weight": 1, "level": 0},
    {"id": 1, "parent_id": 0, "name": "长子", "weight": 2, "level": 1},
    {"id": 2, "parent_id": 0, "name": "次子", "weight": 3, "level": 1},
    {"id": 3, "parent_id": 1, "name": "长孙", "weight": 1, "level": 2},
    {"id": 4, "parent_id": 1, "name": "次孙", "weight": 1, "level": 2},
    {"id": 5, "parent_id": 2, "name": "三孙", "weight": 2, "level": 2},
    {"id": 6, "parent_id": 2, "name": "四孙", "weight": 1, "level": 2},
]

2. 计算节点角度范围 #

编写函数计算每个节点在对应层级的起始和结束角度，同时预留缺角：

import matplotlib.pyplot as plt
from matplotlib.patches import Wedge
import numpy as np

def calculate_angles(tree_data, gap_angle=60):
    # 转换角度为弧度
    total_angle = 2 * np.pi - np.deg2rad(gap_angle)
    # 按层级分组
    level_groups = {}
    for node in tree_data:
        level = node["level"]
        if level not in level_groups:
            level_groups[level] = []
        level_groups[level].append(node)
    
    # 计算根节点角度（根节点占满整个有效区域）
    root = level_groups[0][0]
    root["start_angle"] = np.deg2rad(gap_angle/2)  # 缺角居中
    root["end_angle"] = root["start_angle"] + total_angle
    
    # 递归计算子节点角度
    for level in sorted(level_groups.keys())[1:]:
        parent_level = level - 1
        for parent in level_groups[parent_level]:
            parent_weight = parent["weight"]
            parent_angle_range = parent["end_angle"] - parent["start_angle"]
            # 获取当前父节点的所有子节点
            children = [n for n in level_groups[level] if n["parent_id"] == parent["id"]]
            total_child_weight = sum(c["weight"] for c in children)
            current_start = parent["start_angle"]
            for child in children:
                angle_ratio = child["weight"] / total_child_weight
                child_angle = parent_angle_range * angle_ratio
                child["start_angle"] = current_start
                child["end_angle"] = current_start + child_angle
                current_start = child["end_angle"]
    return tree_data

3. 绘制带缺角的旭日图 #

def draw_sunburst(tree_data, gap_angle=60, ring_width=0.8):
    fig, ax = plt.subplots(figsize=(8,8), subplot_kw={"projection": "polar"})
    ax.set_theta_zero_location("N")  # 角度从顶部开始
    ax.set_theta_direction(-1)  # 顺时针方向
    ax.set_yticklabels([])
    ax.set_xticklabels([])
    
    # 获取最大层级
    max_level = max(n["level"] for n in tree_data)
    
    for node in tree_data:
        level = node["level"]
        # 计算环形的内半径和外半径
        inner_radius = level * ring_width
        outer_radius = (level + 1) * ring_width
        # 创建楔形
        wedge = Wedge((0,0), outer_radius, 
                      np.rad2deg(node["start_angle"]), 
                      np.rad2deg(node["end_angle"]),
                      width=ring_width,
                      facecolor=np.random.rand(3,),  # 随机颜色，可自定义
                      edgecolor="white")
        ax.add_patch(wedge)
        
        # 添加节点标签（在环形中间位置）
        mid_angle = (node["start_angle"] + node["end_angle"]) / 2
        label_radius = inner_radius + ring_width/2
        x = label_radius * np.cos(mid_angle)
        y = label_radius * np.sin(mid_angle)
        ax.text(x, y, node["name"], ha="center", va="center", fontsize=10)
    
    # 绘制缺角的边界线（可选）
    gap_start = np.deg2rad(-gap_angle/2)
    gap_end = np.deg2rad(gap_angle/2)
    for r in np.arange(0, (max_level+1)*ring_width, ring_width):
        ax.plot([gap_start, gap_start], [0, r], color="black", linewidth=1)
        ax.plot([gap_end, gap_end], [0, r], color="black", linewidth=1)
    
    plt.show()

# 执行流程
tree_with_angles = calculate_angles(tree_data)
draw_sunburst(tree_with_angles)

自定义调整 #

• 缺角大小：修改gap_angle参数（默认60°），可调整缺角的宽度
• 环形宽度：修改ring_width参数，控制每层环形的厚度
• 颜色方案：替换facecolor=np.random.rand(3,)为自定义颜色列表，比如根据层级或节点类型设置统一颜色
• 树结构扩展：只需在tree_data中添加更多节点，确保parent_id和level正确即可适配更大的二叉搜索树

内容的提问来源于stack exchange，提问作者Víctor Martínez