如何在Python中生成一级分类以颜色和图例区分、仅显示二级方块的树形图

针对你的需求——生成一级分类用颜色区分、二级分类方块规整填充对应一级区域、仅显示二级标签的静态树形图，最适合的工具是结合squarify与matplotlib实现层级划分，或者调整plotly的配置来达到效果。以下是具体实现：

方案1：squarify + matplotlib 层级树形图（推荐，完全匹配需求） #

这个方法会先按一级分类（city）划分大矩形区域，再在每个区域内填充二级分类（product）的小方块，既保证一级分类的规整性，又能清晰区分二级元素。

代码实现 #

import pandas as pd
import matplotlib.pyplot as plt
import squarify
import seaborn as sns

# 测试数据
product = ["Pins", "Binders", "Paper", "Pins", "Binders", "Paper"]
city = ["Amsterdam", "Amsterdam", "Amsterdam", "Brussels", "Brussels", "Brussels"]
sales = [6, 4, 1, 5, 3, 1]
df = pd.DataFrame(dict(product=product, city=city, sales=sales))

# 1. 计算一级分类（city）的总销售额，用于划分大区域
city_totals = df.groupby('city')['sales'].sum().reset_index()
# 为每个一级分类分配颜色
color_palette = sns.color_palette("deep", n_colors=len(city_totals))
city_colors = dict(zip(city_totals['city'], color_palette))

# 2. 创建画布
fig, ax = plt.subplots(figsize=(12, 8))
ax.set_axis_off()

# 3. 先绘制一级分类的大背景矩形（增强视觉分组）
squarify.plot(
    sizes=city_totals['sales'],
    color=[city_colors[city] for city in city_totals['city']],
    alpha=0.3,  # 背景透明度，突出二级方块
    ax=ax
)

# 4. 为每个一级分类，在对应区域内绘制二级分类的方块
# 获取一级分类的大区域坐标
parent_rects = squarify.squarify(city_totals['sales'], 0, 0, 100, 100)

for i, city_row in city_totals.iterrows():
    # 筛选当前city的二级数据
    subset = df[df['city'] == city_row['city']]
    # 在父区域内计算二级方块的坐标
    child_rects = squarify.squarify(
        subset['sales'],
        parent_rects[i]['x'],
        parent_rects[i]['y'],
        parent_rects[i]['dx'],
        parent_rects[i]['dy']
    )
    # 绘制二级方块（用一级分类的颜色，透明度更高）
    squarify.plot(
        sizes=subset['sales'],
        color=[city_colors[city_row['city']]],
        label=subset['product'],
        ax=ax,
        text_kwargs={'fontsize': 12}
    )

# 添加图例
from matplotlib.patches import Patch
legend_elements = [Patch(facecolor=color, label=city) for city, color in city_colors.items()]
ax.legend(handles=legend_elements, loc='upper right', fontsize=12)

plt.tight_layout()
plt.savefig('treemap_static.png', dpi=300)
plt.show()

效果说明 #

• 一级分类（Amsterdam/Brussels）以不同颜色的大区域区分，区域大小对应总销售额
• 二级分类（Pins/Binders/Paper）的方块完全填充对应一级区域，仅显示二级标签
• 同属一个一级分类的二级元素紧密排列，直观体现一级分类的整体规模

方案2：调整Plotly Treemap 隐藏一级标签 #

如果你更熟悉Plotly，可以通过配置隐藏一级分类的显示，让二级方块看起来填满整个一级区域：

代码实现 #

import pandas as pd
import plotly.express as px
import seaborn as sns

# 测试数据
product = ["Pins", "Binders", "Paper", "Pins", "Binders", "Paper"]
city = ["Amsterdam", "Amsterdam", "Amsterdam", "Brussels", "Brussels", "Brussels"]
sales = [6, 4, 1, 5, 3, 1]
df = pd.DataFrame(dict(product=product, city=city, sales=sales))

fig = px.treemap(
    df,
    path=['city', 'product'],
    values='sales',
    color='city',
    color_discrete_map=dict(zip(df['city'].unique(), sns.color_palette("deep").as_hex()))
)

# 隐藏一级分类的标签和边框，仅显示二级内容
fig.update_traces(
    textinfo='label+value',  # 仅显示二级标签和数值
    marker=dict(cornerradius=3),
    hoverinfo='label+value+parent'  # hover时可查看所属一级分类
)
# 调整布局，优化显示效果
fig.update_layout(
    margin=dict(t=50, l=25, r=25, b=25),
    treemapcolorway=sns.color_palette("deep").as_hex()
)

# 导出为静态图片（需要安装kaleido：pip install kaleido）
fig.write_image('treemap_plotly.png', scale=3)
fig.show()

注意事项 #

• 需要额外安装kaleido包才能导出静态图片
• 虽然视觉上二级方块填满一级区域，但Plotly的treemap本质还是嵌套结构，不过通过隐藏一级标签可以达到你要的效果

选择建议 #

• 如果追求完全的静态层级分组效果，优先选方案1，它完全符合你对一级分类规整性的要求
• 如果习惯Plotly的语法，且可以接受嵌套结构的隐藏处理，选方案2

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