如何为Seaborn Heatmap添加Y轴标签图例色条

要实现给heatmap的Y轴添加对应行标签的色条，我们可以借助Matplotlib的子图布局，在主heatmap旁单独绘制一个颜色映射区域，具体实现如下：

完整代码示例 #

import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from matplotlib.patches import Patch

# 生成测试数据
test = pd.DataFrame.from_dict({
    "A": [np.random.randint(10) for _ in range(10)],
    "B": [np.random.randint(10) for _ in range(10)],
})
labels = ["class1", "class1", "class1",  "class1", "class1", "class2", "class2", "class3", "class3", "class3"]

# 将标签转换为可映射的数值
label_set = list(set(labels))
label_map = {label: idx for idx, label in enumerate(label_set)}
label_nums = [label_map[label] for label in labels]

# 为每个类别指定专属颜色
label_cmap = ListedColormap(["#1f77b4", "#ff7f0e", "#2ca02c"])

# 创建画布与子图：右侧为主heatmap，左侧为标签色条
fig, (ax_label, ax_heat) = plt.subplots(ncols=2, figsize=(6, 5), gridspec_kw={"width_ratios": [0.08, 1]})

# 绘制主heatmap，关闭默认色条避免冲突
sns.heatmap(test, ax=ax_heat, cbar=True)

# 绘制侧边标签色条：垂直展示类别颜色
ax_label.imshow(np.array(label_nums).reshape(-1, 1), cmap=label_cmap, aspect="auto")
# 清理色条的多余元素
ax_label.set_xticks([])
ax_label.set_yticks(np.arange(len(labels)))
ax_label.set_yticklabels(labels)
ax_label.spines[["top", "right", "bottom"]].set_visible(False)

# 添加类别-颜色对应图例
legend_patches = [Patch(facecolor=label_cmap(idx), label=label) for idx, label in enumerate(label_set)]
fig.legend(handles=legend_patches, loc="upper center", bbox_to_anchor=(0.5, 1.05), ncol=3)

plt.tight_layout()
plt.show()

核心逻辑说明 #

• 标签数值化：将字符串标签转换为连续数值，让Matplotlib可以通过颜色映射匹配类别
• 子图布局：通过gridspec_kw控制色条与主图的宽度比例，保证视觉协调
• 色条绘制：用imshow将数值化的标签转为垂直颜色块，同时保留原标签文本作为Y轴刻度
• 图例补充：通过自定义Patch元素创建图例，明确颜色与类别的对应关系

如果需要将色条放在heatmap右侧，只需调整子图的顺序和布局参数即可。

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