matplotlib scatterplot with legend

Actually both linked questions provide a way how to achieve the desired result.

The easiest method is to create as many scatter plots as unique classes exist and give each a single color and legend entry.

import matplotlib.pyplot as plt

x=[1,2,3,4]
y=[5,6,7,8]
classes = [2,4,4,2]
unique = list(set(classes))
colors = [plt.cm.jet(float(i)/max(unique)) for i in unique]
for i, u in enumerate(unique):
    xi = [x[j] for j  in range(len(x)) if classes[j] == u]
    yi = [y[j] for j  in range(len(x)) if classes[j] == u]
    plt.scatter(xi, yi, c=colors[i], label=str(u))
plt.legend()

plt.show()

enter image description here

In case the classes are string labels, the solution would look slightly different, in that you need to get the colors from their index instead of using the classes themselves.

import numpy as np
import matplotlib.pyplot as plt

x=[1,2,3,4]
y=[5,6,7,8]
classes = ['X','Y','Z','X']
unique = np.unique(classes)
colors = [plt.cm.jet(i/float(len(unique)-1)) for i in range(len(unique))]
for i, u in enumerate(unique):
    xi = [x[j] for j  in range(len(x)) if classes[j] == u]
    yi = [y[j] for j  in range(len(x)) if classes[j] == u]
    plt.scatter(xi, yi, c=colors[i], label=str(u))
plt.legend()

plt.show()

enter image description here


Maybe manually filling a table could be useful here. Another idea is using colorbar if your classes are contiguous numbers. I'm showing both approaches in one.

import matplotlib.pyplot as plt
import numpy as np

x=[1,2,3,4,5,6,7]
y=[1,2,3,4,5,6,7]
classes = [2,4,4,2,1,3,5]
cmap = plt.cm.get_cmap("viridis",5)
plt.scatter(x, y, c=classes, label=classes,cmap=cmap,vmin=0.5,vmax=5.5)
plt.colorbar()
unique_classes = list(set(classes))
plt.table(cellText=[[x] for x in unique_classes], loc='lower right',
          colWidths=[0.2],rowColours=cmap(np.array(unique_classes)-1),
         rowLabels=['label%d'%x for x in unique_classes],
          colLabels=['classes'])

enter image description here