Python实现树状图分支按分类向上着色的方法求助

要实现将叶子标签的着色规则延伸至分支，SciPy的hierarchy.dendrogram没有直接提供该功能，需要手动解析树状图结构并修改分支颜色。核心思路是：根据叶子的类别，向上遍历分支节点，若当前节点下所有叶子属于同一类别，则为该分支上色，直到遇到包含多类别的节点为止。

修改后的完整代码 #

import pandas as pd
from scipy.cluster import hierarchy
import matplotlib.pyplot as plt

# 假设dist_matrix和df是已定义的输入数据
linkage_matrix = hierarchy.ward(dist_matrix)

# 绘制树状图并获取返回的结构信息
dendro_info = hierarchy.dendrogram(
    linkage_matrix,
    labels=df.index,
    leaf_rotation=0,
    orientation="left",
    color_threshold=None,  # 先关闭默认分支着色，避免冲突
    above_threshold_color='grey',
    leaf_font_size=10
)

# 处理类别与颜色映射
df['Location'] = pd.Categorical(df['Location'])
my_palette = plt.colormaps["Accent"]
# 根据树状图的叶子顺序，匹配对应类别的颜色
leaf_order = dendro_info['leaves']
leaf_colors = [my_palette(df['Location'].cat.codes.iloc[i]) for i in leaf_order]

# 获取当前绘图轴对象
ax = plt.gca()

# 定义函数：递归获取指定节点下的所有原始叶子样本索引
def get_node_leaves(node_idx, linkage_matrix):
    leaves = []
    stack = [node_idx]
    n_samples = len(linkage_matrix) + 1  # 原始样本数 = 链接矩阵行数 + 1
    while stack:
        idx = stack.pop()
        if idx < n_samples:
            leaves.append(idx)
        else:
            # 链接矩阵中，合并节点从n_samples开始计数，每行存储左右子节点
            left_node = int(linkage_matrix[idx - n_samples, 0])
            right_node = int(linkage_matrix[idx - n_samples, 1])
            stack.append(right_node)
            stack.append(left_node)
    return leaves

# 遍历所有分支，判断并设置颜色
n_samples = len(df)
for i, (icoord, dcoord) in enumerate(zip(dendro_info['icoord'], dendro_info['dcoord'])):
    # 当前分支对应的合并节点索引
    node_idx = n_samples + i
    # 获取该节点下所有叶子样本
    node_leaves = get_node_leaves(node_idx, linkage_matrix)
    # 检查这些叶子是否属于同一类别
    node_categories = df['Location'].cat.codes.iloc[node_leaves]
    if len(node_categories.unique()) == 1:
        # 所有叶子类别一致，设置分支颜色
        branch_color = my_palette(node_categories.iloc[0])
        ax.plot(icoord, dcoord, color=branch_color)

# 设置叶子标签颜色
xlbls = ax.get_ymajorticklabels()
for lbl, color in zip(xlbls, leaf_colors):
    lbl.set_color(color)

plt.figure(figsize=(10,15))
plt.show()

关键步骤说明 #

• 获取树状图结构：调用hierarchy.dendrogram时保存返回的dendro_info，其中包含分支坐标、叶子排列顺序等核心数据。
• 叶子颜色映射：根据树状图实际的叶子排列顺序，匹配对应的类别颜色，确保标签颜色与叶子位置对应。
• 节点递归遍历：通过get_node_leaves函数，递归拆解每个合并节点，获取其包含的所有原始叶子样本。
• 分支着色判断：若节点下所有叶子属于同一类别，则为该分支设置对应颜色；若包含多类别，则保留默认的灰色。

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