Is there a quick way to set the legend of a scatterplot without looping?

Question

I have 3 arrays like so:

x = [1,2,3,4,5]
y = [5,4,3,2,1]
c = [0,0,1,0,1]

plt.figure(figsize = (12,9))

plt.scatter(x = x, y = y, c = c)

plt.legend(['0', '1'])

I'm able to generate a scatterplot like this:

But what I want is for it to differentiate the colors between 0 and 1.

The solution here does a for loop over the classes to achieve the desired result. I've also tried iterating over the plt.scatter() object but it's not iterable.

Is there some sort of simple solution out there? preferably no loops and only about 1 line of code?

Answer 1

No yet. But Matplotlib will soon release a version of scatter with a straightforward scatter + legend with no need for multiple calls.

Note that you can also use the Seaborn scatterplot function that you use to generate a scatterplot with labels in one single line (see below from the Seaborn documentation):

sns.scatterplot(data=tips, x="total_bill", y="tip", hue="time")

Answer 2

You can loop and plot all different sets of data one by one with appropriate label. Here I am plotting the red dots then the green ones.

import matplotlib.pyplot as plt
x = [1,2,3,4,5]
y = [5,4,3,2,1]
c = [0,0,1,0,1]

unique = list(set(c))
colors = ['red','green']
for i, u in enumerate(unique):
    xi = [x[j] for j  in range(len(x)) if c[j] == u]
    yi = [y[j] for j  in range(len(x)) if c[j] == u]
    plt.scatter(xi, yi, c=colors[i], label=str(u))
plt.legend()

plt.show()