This is k-means implementation using Python (numpy). I believe there is room for improvement when it comes to computing distances (given I'm using a list comprehension, maybe I could also pack it in a numpy operation) and to compute the centroids using label-wise means (which I think also may be packed in a numpy operation).
import numpy as np def k_means(data, k=2, max_iter=100): """Assigns data points into clusters using the k-means algorithm. Parameters ---------- data : ndarray A 2D array containing data points to be clustered. k : int, optional Number of clusters (default = 2). max_iter : int, optional Number of maximum iterations Returns ------- labels : ndarray A 1D array of labels for their respective input data points. """ # data_max/data_min : array containing column-wise maximum/minimum values data_max = np.max(data, axis=0) data_min = np.min(data, axis=0) n_samples = data.shape n_features = data.shape # labels : array containing labels for data points, randomly initialized labels = np.random.randint(low=0, high=k, size=n_samples) # centroids : 2D containing centroids for the k-means algorithm # randomly initialized s.t. data_min <= centroid < data_max centroids = np.random.uniform(low=0., high=1., size=(k, n_features)) centroids = centroids * (data_max - data_min) + data_min # k-means algorithm for i in range(max_iter): # distances : between datapoints and centroids distances = np.array( [np.linalg.norm(data - c, axis=1) for c in centroids]) # new_labels : computed by finding centroid with minimal distance new_labels = np.argmin(distances, axis=0) if (labels == new_labels).all(): # labels unchanged labels = new_labels print('Labels unchanged ! Terminating k-means.') break else: # labels changed # difference : percentage of changed labels difference = np.mean(labels != new_labels) print('%4f%% labels changed' % (difference * 100)) labels = new_labels for c in range(k): # computing centroids by taking the mean over associated data points centroids[c] = np.mean(data[labels == c], axis=0) return labels