如何用Python生成一级为颜色图例、仅显示二级区块的Treemap

方案一：squarify + Matplotlib（手动分组布局） #

该方案可让同一一级分类的二级区块聚集排布，直观体现一级组的整体大小，仅显示二级标签，颜色按一级分类区分，支持导出静态图片。

代码示例 #

import squarify
import matplotlib.pyplot as plt
import pandas as pd
from matplotlib.patches import Patch

# 1. 准备示例数据
data = pd.DataFrame({
    "category1": ["A", "A", "A", "B", "B", "C", "C", "C", "C"],
    "category2": ["A1", "A2", "A3", "B1", "B2", "C1", "C2", "C3", "C4"],
    "value": [30, 20, 10, 25, 15, 10, 15, 20, 5]
})

# 2. 按一级分类分组计算总大小并排序（优化布局）
grouped = data.groupby("category1").agg(total=("value", "sum")).sort_values("total", ascending=False).reset_index()

# 3. 定义一级分类颜色映射
color_map = {"A": "#FF6B6B", "B": "#4ECDC4", "C": "#45B7D1"}

# 4. 初始化画布并绘制一级组占位区块
fig, ax = plt.subplots(figsize=(10, 8))
squarify.plot(sizes=grouped["total"], color=[color_map[cat] for cat in grouped["category1"]], ax=ax)
ax.set_axis_off()

# 5. 在每个一级组区域内填充二级区块并添加标签
patch_idx = len(grouped)  # 记录二级区块的起始索引
for idx, row in grouped.iterrows():
    group_data = data[data["category1"] == row["category1"]]
    # 获取当前一级组的区域边界
    parent_patch = ax.patches[idx]
    x, y, dx, dy = parent_patch.get_x(), parent_patch.get_y(), parent_patch.get_width(), parent_patch.get_height()
    # 计算二级区块相对大小
    relative_sizes = group_data["value"] / row["total"] * dx * dy
    # 绘制二级区块
    squarify.plot(sizes=relative_sizes, color=[color_map[row["category1"]]]*len(group_data),
                  ax=ax, x=x, y=y, dx=dx, dy=dy)
    # 添加二级标签
    for i, (_, item) in enumerate(group_data.iterrows()):
        child_patch = ax.patches[patch_idx]
        rx, ry, rdx, rdy = child_patch.get_x(), child_patch.get_y(), child_patch.get_width(), child_patch.get_height()
        ax.text(rx + rdx/2, ry + rdy/2, item["category2"], ha="center", va="center", fontsize=10)
        patch_idx += 1

# 6. 添加图例
legend_elements = [Patch(facecolor=color_map[cat], label=cat) for cat in grouped["category1"]]
ax.legend(handles=legend_elements, loc="upper right", bbox_to_anchor=(1.15, 1))

# 7. 导出静态图片
plt.savefig("treemap_matplotlib.png", bbox_inches="tight", dpi=300)
plt.show()

核心逻辑 #

先绘制一级组的占位区块确定整体大小比例，再在每个一级区域内填充二级区块，保证同组二级区块聚集；颜色由一级分类统一控制，仅显示二级标签。

方案二：Plotly 调整参数（隐藏一级区域） #

Plotly默认会显示一级父级区域，可通过参数隐藏父级的标签、边框，仅展示二级区块，同时按一级分类着色。

代码示例 #

import plotly.express as px
import pandas as pd

# 准备示例数据
data = pd.DataFrame({
    "category1": ["A", "A", "A", "B", "B", "C", "C", "C", "C"],
    "category2": ["A1", "A2", "A3", "B1", "B2", "C1", "C2", "C3", "C4"],
    "value": [30, 20, 10, 25, 15, 10, 15, 20, 5]
})

# 生成Treemap并配置参数
fig = px.treemap(data, path=["category1", "category2"], values="value",
                 color="category1",
                 color_discrete_map={"A": "#FF6B6B", "B": "#4ECDC4", "C": "#45B7D1"})

# 隐藏一级区域的显示元素
fig.update_traces(
    textinfo="label",  # 仅显示二级标签
    pathbar_visible=False,  # 隐藏路径栏（避免显示一级分类路径）
    marker=dict(cornerradius=3)  # 可选：设置区块圆角优化视觉效果
)

# 调整布局边距
fig.update_layout(margin=dict(t=50, l=25, r=25, b=25))

# 导出静态图片（需提前安装kaleido：pip install kaleido）
fig.write_image("treemap_plotly.png", width=1000, height=800, scale=3)
fig.show()

核心逻辑 #

通过path设置层级关系，用textinfo和pathbar_visible参数隐藏一级分类的显示；二级区块继承一级分类的颜色，整体布局保留一级组的大小比例。

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