sklearnでクラスタリング(その2)

プロットするプログラムhello_scatterもつけてみた。

import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans

def hello_kmeans_demo():
    X = np.array([[1, 2], [1, 4], [1, 0],[10, 2], [10, 4], [10, 0]])
    """
    array([[ 1,  2],
       [ 1,  4],
       [ 1,  0],
       [10,  2],
       [10,  4],
       [10,  0]])
    X.shapeは、(6,2)
    """
    kmeans = KMeans(n_clusters=2, random_state=0).fit(X)
　#クラスタリングの結果のラベルが手に入る
    Y=np.array(kmeans.labels_)

    #print("cluster_centers_",kmeans.cluster_centers_)
    #print("inertia_",kmeans.inertia_)
    #print("kmeans.n_iter_",kmeans.n_iter_)

　#プロット関数呼び出し
    hello_scatter(X,Y)

def hello_scatter(np_data,np_target=None,dim_x=0,dim_y=1):
    """
    指定された次元で、scatterプロットする。
    targetのラベルが0ならば赤、1ならば緑、2ならば青、それ以外は黒でプロット。
    """
    color_list=["black"]*np_data.shape[0]#shapeで(行,列)が返ってくるので行数分のリストを作成。
    for i in np.where(np_target==0)[0]:
        color_list[i]="red"
    for i in np.where(np_target==1)[0]:
        color_list[i]="green"
    for i in np.where(np_target==2)[0]:
        color_list[i]="blue"

    #arrayは、[行,列]で指定
    plt.scatter(x=np_data[:,dim_x], y=np_data[:,dim_y], c=color_list, marker=".")
    plt.grid(True)
    plt.show()

if __name__ == "__main__":
    hello_kmeans_demo()