如何基于pandas DataFrame绘制带自定义颜色对应图例的scatter散点图

问题原因 #

直接向plt.scatter()的c参数传入颜色序列时，matplotlib不会自动为每个颜色绑定标签生成图例，需要手动构造图例映射关系。

修复代码 #

你可以通过matplotlib.patches.Patch手动创建图例句柄，绑定你预设的palette颜色映射即可，完整可运行代码如下：

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Patch

x = np.array([1,2,3,4,5,6,7,8,9,10])
y = 2*x

df = pd.DataFrame()
palette = {1: "blue", 2:"orange", 3:"green", 4:"red", 5:"purple", 6:"brown", 7:"pink", 8:"gray", 9:"olive", 10:"cyan"}
df["first"] = x
df["second"] = y
df["third"] = df["first"].apply(lambda x: palette[x])

plt.scatter(df["first"], df["second"], c=df["third"])

# 手动构造图例句柄
legend_handles = [Patch(color=color, label=str(num)) for num, color in palette.items()]
plt.legend(handles=legend_handles, title="数值对应颜色")

plt.show()

适配实际场景说明 #

如果实际数据中first列的取值没有覆盖palette的所有键，你可以过滤只保留数据中存在的数值来生成图例，避免展示无用条目：

# 仅保留数据中存在的数值对应的图例
exist_nums = df["first"].unique()
legend_handles = [Patch(color=palette[num], label=str(num)) for num in exist_nums]

内容的提问来源于stack exchange，提问作者Josh.h