Question

I’d like to add a colored legend to a Matplotlib scatterplot. Here’s my code:

xs = [1, 2, 1, 4, 3, 2]
ys = [1, 3, 2, 2, 3, 1]
labels = [1, 1, 0, 2, 1, 3]

label_dict = {0: 'r', 1: 'k', 2: 'b', 3: 'g'}
legend_dict = {0: 'foo', 1: 'bar', 2: 'baz', 3: 'biff'}

for x, y, label in zip(xs, ys, labels):
    plt.scatter(x, y, c=label_dict.get(label), label=legend_dict.get(label))

plt.legend()
plt.show()

How can I get the legend to only display one label for each color instead of a label for each point?

Answer 1

You can keep track of which labels you’ve seen:

import pylab as plt

xs = [1, 2, 1, 4, 3, 2]
ys = [1, 3, 2, 2, 3, 1]
labels = [1, 1, 0, 2, 1, 3]

label_dict = {0: 'r', 1: 'k', 2: 'b', 3: 'g'}
legend_dict = {0: 'foo', 1: 'bar', 2: 'baz', 3: 'biff'}

seen = set()
for x, y, label in zip(xs, ys, labels):
    if label not in seen:
        plt.scatter(x, y, c=label_dict.get(label), label=legend_dict.get(label))
    else:
        plt.scatter(x, y, c=label_dict.get(label))
    seen.add(label)

plt.legend()
plt.show()

The if/else clause could be condensed into 1 line if you’d rather:

seen = set()
for x, y, label in zip(xs, ys, labels):
    plt.scatter(x, y, c=label_dict.get(label), label=legend_dict.get(label) if label not in seen else None)
    seen.add(label)

---

I think that I would personally prefer to keep the data grouped. In other words, I’d probably store all of the data with the same label together, then you only need to issue one plot command per label type:

import numpy as np
import pylab as plt

xs = [1, 2, 1, 4, 3, 2]
ys = [1, 3, 2, 2, 3, 1]
labels = [1, 1, 0, 2, 1, 3]

xs = np.array(xs)
ys = np.array(ys)
labels = np.array(labels)

labels_masks =( (x,(labels == x)) for x in set(labels))
data_dict = dict( (lbl,(xs[mask],ys[mask])) for lbl,mask in labels_masks )


label_dict = {0: 'r', 1: 'k', 2: 'b', 3: 'g'}
legend_dict = {0: 'foo', 1: 'bar', 2: 'baz', 3: 'biff'}

for label,data in data_dict.items():
    x,y = data
    plt.scatter(x,y,c=label_dict.get(label),label=legend_dict.get(label))

plt.legend()
plt.show()