python - 描绘没有重复的传奇

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