Python散布図第３変数をラベルとしてつける方法

Question

#k-means法でクラスタリング
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
n_clusters = 3
#モデルの定義から
model = KMeans(n_clusters = n_clusters)
pred = model.fit_predict(df_test)
#print(pred)
centers = model.cluster_centers_

plt.scatter(df['alcohol'],df['color_intensity'])

plt.plot(centers[:,0],centers[:,1], 'o', color='red')

このグラフをクラスターで分類した種類別に色分けしたいです、
pred変数には、178までの分類情報が含んでるのですが、(データフレームの長さも178)

plt.scatter(df['alcohol'],df['color_intensity'],label=pred)

としてもうまくラベルできません。どのようにしたらできるか教えて欲しいです。

Answer 1

for p in set(pred):
    idx = pred == p
    plt.scatter(df.loc[idx, 'alcohol'], df.loc[idx, 'color_intensity'], label=str(p))

としてみてください。また、必要に応じてplt.legend()で凡例を追加できます。

また、単に色分けして散布図を描くだけなら

plt.scatter(df['alcohol'], df['color_intensity'], c=pred)

でもできますが、この場合は一種類のラベルになるため、上述の凡例が出せない他色々な制約があります。

Comments

[chgrios]

解決しました！とても親切にありがとうございます！