描绘没有重复的图例

Question

我想在Matplotlib散点图中添加彩色图例。这是我的代码：

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

如何让图例只为每种颜色显示一个标签而不是每个点的标签？

Answer 1

您可以跟踪您看过的标签：

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

如果您愿意，可以将if / else子句压缩为1行：

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)

---

我认为我个人更愿意将数据分组。换句话说，我可能将所有具有相同标签的数据存储在一起，然后您只需要为每个标签类型发出一个plot命令：

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