如何修改Matplotlib散点图legend_elements()的标签为自定义分类值

嘿，我看你遇到了图例标签显示数字而不是name列的aa/bb/cc的问题，这是因为默认的legend_elements()不知道你的颜色映射对应着name的分类。咱们来改改代码，让图例精准对应上你的分类名称，同时保持颜色不变：

修改后的完整代码 #

import matplotlib.pyplot as plt
import pandas as pd

# 构造你的DataFrame（如果已经存在可以忽略这段）
df = pd.DataFrame({
    'age': [12,23,34,32,23,12],
    'weight': [100,120,121,134,133,112],
    'score': [12,12,13,10,11,19],
    'height': [501,502,499,499,498,503],
    'name': ['aa','bb','bb','cc','cc','aa']
})

fig, axes = plt.subplots(2,2,figsize=(12,8))
properties = ['age','weight','score','height']
indexes = df.index.tolist()

# 先给name列的分类做个映射：把每个唯一名称对应到一个数字编码
unique_names = df['name'].unique()
name_code_map = {name: idx for idx, name in enumerate(unique_names)}
# 生成每个数据点对应的编码数组
point_codes = df['name'].map(name_code_map).values

for counter in range(4):
    i = counter % 2
    j = counter // 2  # 用整数除法替代math.floor，更简洁
    # 绘制散点图时，用point_codes作为颜色参数，保证同name同色
    scatter = axes[i,j].scatter(indexes, df[properties[counter]], c=point_codes)
    axes[i,j].set_title(properties[counter])
    # 自定义图例：用unique_names作为标签，对应散点的颜色手柄
    legend_handles = scatter.legend_elements()[0]
    axes[i,j].legend(handles=legend_handles, labels=unique_names)

plt.tight_layout()
plt.show()

核心修改说明 #

• 分类编码映射：先提取name的唯一值，给每个名称分配一个数字，这样c=point_codes就能让相同name的点拥有相同颜色。
• 自定义图例：不再使用legend_elements()默认生成的数字标签，而是手动指定labels=unique_names，同时用scatter.legend_elements()[0]拿到对应颜色的图例手柄，这样图例就会显示aa、bb、cc啦。
• 顺手把math.floor(counter/2)换成了Python的整数除法counter//2，代码更简洁易读。

这样修改后，每个子图的图例都会准确对应name列的分类，颜色也会和每个分类保持一致哦～

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