Python pyplot 图例散点图

Question

我有 5 个类和一些要绘制的特征。这是代码

x_pts = X_test.iloc[:,col_1]
y_pts = X_test.iloc[:,col_2]
color_seq = y_test
plt.scatter(x_pts, y_pts, c=color_seq, cmap='viridis')
plt.xlabel(X_test.columns[col_1])
plt.ylabel(X_test.columns[col_2])
plt.show()

结果如下图

我现在想要每种颜色的图例（例如，黄色 = 'class a' ，蓝色 = 'class b'，...） 我能找到的唯一文档是人们以不同的方式绘制每种颜色，这在我的具体情况下非常困难。有没有简单的显示图例的方法like the example here

Answer 1

希望这会有所帮助：

x = [1, 3, 4, 6, 7, 9]
y = [0, 0, 5, 8, 8, 8]
labels = ['A', 'B', 'C']
colors = [0, 0, 1, 2, 2, 2]
scatter = plt.scatter(x, y,c=colors, cmap='viridis')
plt.legend(handles=scatter.legend_elements()[0], labels=labels)
plt.show()

输出：