K Means clustering

kmeans.py

import numpy as np, matplotlib.pyplot as plt

def KMeans(X, K, num_steps=100):
    L = len(X)
    cluster_assignments=np.concatenate([np.full((L//K,), k) for k in range(K-1)]+[np.full((L-(K-1)*(L//K),), K-1)])
    centroids = np.empty((K,*X.shape[1:]))
    for step in range(num_steps):
        # E step: centroid calculation
        for k in range(K):
            centroids[k] = X[cluster_assignments==k].mean(axis=0)
        # M step: cluster reassignment
        distances = np.array([[np.linalg.norm(x-c) for c in centroids] for x in X])
        cluster_assignments = distances.argmin(axis=1)
    print("K=", K)
    print(cluster_assignments)

    plt.clf()
    for k in range(K):
        plt.scatter(X[cluster_assignments==k,0], X[cluster_assignments==k,1], label="C={}".format(k))
    plt.legend()
    plt.savefig("kmeans_K_{}.jpg".format(K))

def main():
    # {[−2,−1]T ,[−3,−2]T ,[0,−1]T ,[−1,0]T ,[2,3]T ,[−1,−2]T ,[3,2]T ,[3,3]T ,[1,1]T ,[2,2]T }
    X = np.array([[-2,-1], [-3,-2], [0,-1], [-1,0], [2,3], [-1,-2], [3,2], [3,3], \
                  [1,1], [2,2]])
    KMeans(X, 2)
    KMeans(X, 3)
    KMeans(X, 4)
    KMeans(X, 5)
    KMeans(X, 6)
    KMeans(X, 8)
    KMeans(X, 9)
    KMeans(X, 10)

main()