Legends without repeats

Question

I would like to add a color legend to the Matplotlib scatter chart. Here is 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 make a legend display only one label for each color instead of a label for each point?

Answer 1

You can track which tags you saw:

 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 can be compressed into 1 line if you prefer:

 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 personally would prefer to keep the data grouped. In other words, I would probably save all the data with the same label together, then you only need to issue one plot command for each type of label:

 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()