【Python・matplotlib】クラスタリング結果のクラスタ番号を凡例表示

Question

scikit-learnで行ったクラスタリング（k-MeanやGMM）の結果から得たクラスタIDを元
のデータフレームに結合して、散布図化しています。

Name	value1	value2	ClusterID
AAA	100	200	0
BBB	200	100	0
CCC	200	500	1
DDD	300	500	1
EEE	150	800	2

散布図の色からクラスタIDを判断したいので、

plt.scatter(df['value1'], df['value2'], c=df['ClusterID'])

上記で表示させるグラフの凡例としてClusterIDを表示したいのですが、どのようにすればよいでしょうか？

plt.scatterの中で、label=""を指定するのだと思うのですが…上手くいきません。

ご教示頂ければ幸いです。
よろしくお願いいたします。

Answer 1

ループを使って良いならこんな感じでかけます。

import pandas as pd
import matplotlib.pyplot as plt

df = pd.DataFrame([['AAA', 100, 200, 0],
                   ['BBB', 200, 100, 0],
                   ['CCC', 200, 500, 1],
                   ['DDD', 300, 500, 1],
                   ['EEE', 150, 800, 2]],
                  columns=['Name','value1','value2','ClusterID'])

ClusterName = ['Cluster0','Cluster1','Cluster2']
ax = plt.subplot()
for id, d in df.groupby('ClusterID'):
    ax.scatter(x=d.value1, y=d.value2, label=ClusterName[id])
ax.legend()
plt.show()

Comments

[TNF_1023]

おかげ様でうまく出来ました。
ありがとうございます。