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"用于将坐标轴外的图例一起纳入导出范围。
备注:内容仅供参考。
