setting a legend matching the colours in pyplot.scatter

Question

Suppose my data is organized in the following way:

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]
label         


        

Answer 1

If you want to use a colormap you can create a legend entry for each unique entry in the colors list as shown below. This approach works well for any number of values. The legend handles are the markers of a plot, such that they match with the scatter points.

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()