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

各ポイントのラベルではなく、各色のラベルを1つだけ表示するように凡例を取得するにはどうすればよいですか？

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つまり、おそらく同じラベルのすべてのデータを一緒に保存するので、ラベルタイプごとに1つのコマンドを発行するだけで済みます。

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