PAYATHON 2026

Python treemap 如何仅显示第二层级,用颜色和图例表示第一层级?

数据阿明

可以先用 squarify 按城市销售总额划分矩形,再在每个城市的矩形内按产品销售额划分。绘图时用 Matplotlib 只绘制产品矩形,将城市名称放入图例。

为什么现有代码没有得到目标效果

原来的 Plotly 代码使用 path=['city', 'product'],已经表达了城市和产品两个层级。当前呈现方式中,父级占用了标题区域,层级之间也有留白。

原来的 squarify.plot() 只接收了一份包含六条记录的 sizes。color 只负责给矩形着色,不会告诉布局算法哪些产品必须放在同一个城市矩形中。因此,同色矩形不一定组成一个完整的矩形区域。

要保持城市分组的矩形形状,需要在布局时处理层级关系。只调整颜色或输入顺序无法保证这一点。

完整实现

示例使用原问题中的六条数据,城市和组内产品都按销售额降序布局。绘图区没有父级标签栏,也没有为父级预留内边距。

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

# 原问题中的测试数据
df = pd.DataFrame({
    "product": ["Pins", "Binders", "Paper",
                "Pins", "Binders", "Paper"],
    "city": ["Amsterdam", "Amsterdam", "Amsterdam",
             "Brussels", "Brussels", "Brussels"],
    "sales": [6, 4, 1, 5, 3, 1],
})

# 本示例要求 sales 为有限的正数。
city_totals = (
    df.groupby("city", sort=False)["sales"]
    .sum()
    .sort_values(ascending=False)
)

city_colors = {
    "Amsterdam": "#4C72B0",
    "Brussels": "#DD8452",
}

width, height = 100, 60

# 第一层:根据城市销售总额分配矩形。
city_sizes = squarify.normalize_sizes(
    city_totals.tolist(), width, height
)
city_rects = squarify.squarify(
    city_sizes, 0, 0, width, height
)

fig, ax = plt.subplots(figsize=(12, 8))

for city, parent in zip(city_totals.index, city_rects):
    # 同一城市内,相同产品的记录合并后参与布局。
    products = (
        df.loc[df["city"].eq(city)]
        .groupby("product", sort=False)["sales"]
        .sum()
        .sort_values(ascending=False)
    )

    # 第二层:在当前城市矩形内划分产品矩形。
    product_sizes = squarify.normalize_sizes(
        products.tolist(), parent["dx"], parent["dy"]
    )
    product_rects = squarify.squarify(
        product_sizes,
        parent["x"],
        parent["y"],
        parent["dx"],
        parent["dy"],
    )

    for product, rect in zip(products.index, product_rects):
        patch = Rectangle(
            (rect["x"], rect["y"]),
            rect["dx"],
            rect["dy"],
            facecolor=city_colors[city],
            edgecolor="white",
            linewidth=1,
        )
        ax.add_patch(patch)

        label = ax.text(
            rect["x"] + rect["dx"] / 2,
            rect["y"] + rect["dy"] / 2,
            product,
            ha="center",
            va="center",
            color="white",
            fontsize=12,
            clip_on=True,
        )
        label.set_clip_path(patch)

ax.set_xlim(0, width)
ax.set_ylim(0, height)
ax.set_aspect("equal")
ax.set_axis_off()

legend_handles = [
    Patch(facecolor=city_colors[city], label=city)
    for city in city_totals.index
]
ax.legend(
    handles=legend_handles,
    loc="upper center",
    bbox_to_anchor=(0.5, -0.03),
    ncol=len(legend_handles),
    frameon=False,
)

fig.tight_layout()
fig.savefig("sales_treemap.png", dpi=300, bbox_inches="tight")
plt.show()

白色描边用于区分产品,矩形布局本身没有预留间隔。如果希望同色区域之间也看不到分隔线,可以将 linewidth 设为 0。

为什么两次布局仍能保持面积比例

第二次调用 normalize_sizes() 时,传入的是当前城市矩形的宽高:

product_sizes = squarify.normalize_sizes(
    products.tolist(), parent["dx"], parent["dy"]
)

这样,产品面积的总和等于所属城市矩形的面积。

示例数据中,Amsterdam 的销售额合计为 11,Brussels 为 9,总计为 20。两个城市分别占总面积的 55% 和 45%。

城市中某个产品的面积占比为:

[ \frac{\text{城市销售额}}{\text{总销售额}} \times \frac{\text{产品在该城市的销售额}}{\text{城市销售额}}

\frac{\text{产品在该城市的销售额}}{\text{总销售额}} ]

因此,城市分组保持矩形,每个产品的面积也与销售额成比例。

使用时需要注意什么

  • 数值条件:代码假定 sales 全部为有限正数。零值没有可见面积,可以在布局前过滤;负数和缺失值应先按业务含义处理。
  • 产品归属:代码按城市分别汇总产品。Amsterdam 的 Pins 和 Brussels 的 Pins 是两个叶子节点,不会跨城市合并。
  • 小矩形标签:产品数量增多或数值差异很大时,小矩形可能放不下文字。示例将文字裁剪到矩形范围内;实际使用时也可以按面积阈值隐藏标签。
  • 颜色映射:示例只为输入中的两个城市定义了颜色。增加城市时,需要同步扩展 city_colors。
  • 静态导出:示例保存 PNG。需要矢量图片时,可将保存文件名改为 sales_treemap.svg 或 sales_treemap.pdf;bbox_inches="tight" 用于将坐标轴外的图例一起纳入导出范围。

备注:内容仅供参考。