Search code examples
pythonmatplotlibscatter-plot

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

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: enter image description here

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?

Solution

  • 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")