How to add legend based on color in scatterplot [matplotlib]

Necronet Published at Dev

Necronet

I have been trying to add legend based on the color in a scatter plot in matplotlib. I have seen several example but they assume that each plot will have a label associated to it, and I haven't figure out how to do it without the "label" property and only base it on color. Currently I have the following code

plt.subplot(121)
plt.gca().set_title('Female')
survived = female_data[:,1]
pclass = female_data[:,2]
age = female_data[:,5]
label_color = ['r' if i==0 else 'b' for i in survived]
axes = plt.gca()
axes.get_xaxis().set_visible(False)
plt.scatter(age, pclass, color=label_color)

Which render something like this, I want red and blue to have a single label in the plot. Any idea?

Mark Mikofski

plot with legend

https://colab.research.google.com/notebook#fileId=1RtdwW8ztZbBUpd67RWffKRVt0ae0dNkr

import matplotlib.pyplot as plt
import numpy as np

ages = np.array([1, 2, 1, 3, 1, 2, 2, 1, 1, 2])
pclass = np.arange(10)
survived = np.array([
    True, True, True, False, False,
    False, True, False, True, False
])
plt.plot(
    pclass[survived], ages[survived], 'bo',
    pclass[~survived], ages[~survived], 'ro'
)
plt.legend(['survived', 'died'])

Hopefully you can see my colab notebook or the image. Is this what you wanted? I don't think you want scatter.

Collected from the Internet

Please contact [email protected] to delete if infringement.