matplotlibで散布図に凡例をつける

散布図を描くにはscatterを用いる。scatterとlegendを組み合わせて使う場合の例はこちらにある。

Scatter plots with a legend — Matplotlib 3.1.2 documentation

これによると以下の2つの方法がある

  1. 凡例としてつけたいlabel毎にscatterを呼ぶ
  2. x, y, cにそれぞれリストを渡しscatterを一度だけ呼ぶ

2.の方法では凡例のラベルとしてcolorで指定した数字が入ることになる。scatterにはlabelも引数として渡せるため、ラベルのリストも作って渡してみたが反映されなかった。 そのため凡例として文字列をつけたい場合は1.の方法でやる必要がある。

データがx, y, labelのリスト(data_listとする)として保持されてるとすると1.の場合はこんな感じ

data = {}
for x, y, label in data_list:
    if label not in data:
        data[label] = {'x':[], 'y':[]}

    data[label]['x'].append(x)
    data[label]['y'].append(y)

for label, item in sorted(data.items(), key=lambda x:x[0]):
    plt.scatter(item['x'], item['y'], label=label)

ちなみにdataに対してsortedをかけているのは凡例のラベルを降順に並べ替えるため。