python - 设置与 pyplot.scatter 中的颜色匹配的图例

Question

假设我的数据按以下方式组织:

x_values = [6.2, 3.6, 7.3, 3.2, 2.7]
y_values = [1.5, 3.2, 5.4, 3.1, 2.8]
colours = [1, 1, 0, 1, -1]
labels = ["a", "a", "b", "a", "c"]

我想用这个制作散点图:

axis = plt.gca()
axis.scatter(x_values, y_values, c=colours)

我想要一个包含 3 个类别的图例:“a”、“b”和“c”。

鉴于此列表中的类别与 colours 列表中的点的顺序相匹配，我可以使用 labels 列表来制作此图例吗？

我需要为每个类别单独运行scatter命令吗？

Answer 1

如果您想使用颜色图，您可以为 colors 列表中的每个唯一条目创建一个图例条目，如下所示。这种方法适用于任意数量的值。图例 handle 是绘图的标记，以便它们与散点相匹配。

import matplotlib.pyplot as plt

x_values = [6.2, 3.6, 7.3, 3.2, 2.7]
y_values = [1.5, 3.2, 5.4, 3.1, 2.8]
colors = [1, 1, 0, 1, -1]
labels = ["a", "a", "b", "a", "c"]
clset = set(zip(colors, labels))

ax = plt.gca()
sc = ax.scatter(x_values, y_values, c=colors, cmap="brg")

handles = [plt.plot([],color=sc.get_cmap()(sc.norm(c)),ls="", marker="o")[0] for c,l in clset ]
labels = [l for c,l in clset]
ax.legend(handles, labels)

plt.show()