Data Science K-means Clustering
K-means clustering groups data by iteratively assigning points to the nearest of k centroids and recomputing those centroids until they stop moving, requiring you to choose k in advance.
Clustering Around Centroids
K-means represents each cluster by a single point, its centroid, which is the average position of all the points currently assigned to that cluster. Every data point is grouped with whichever centroid it is closest to, and the algorithm alternates between assigning points and recalculating centroids until the assignments stop changing.
The K-means Algorithm Step by Step
- Choose the number of clusters, k, before running the algorithm.
- Initialize k centroids, typically by picking k data points at random (or using the smarter k-means++ strategy).
- Assign every data point to its nearest centroid, usually measured by Euclidean distance.
- Recompute each centroid as the mean position of all points now assigned to it.
- Repeat the assign-and-recompute steps until centroids stop moving (or a maximum number of iterations is reached).
Unlike hierarchical clustering, where you can build the merge tree first and decide on a cluster count afterward, k-means needs the number of clusters, k, fixed before it starts, because the algorithm has to initialize exactly that many centroids on the first step.
Example
from sklearn.cluster import KMeans
import numpy as np
X = np.array([
[1, 2], [1.5, 1.8], [1.2, 2.3],
[8, 8], [8.5, 7.8], [9, 8.2],
])
kmeans = KMeans(n_clusters=2, n_init=10, random_state=0)
kmeans.fit(X)
print('Cluster labels:', kmeans.labels_)
print('Centroids:\n', kmeans.cluster_centers_.round(2))
print('Inertia:', round(kmeans.inertia_, 2))Since k is not known ahead of time in most real problems, a common technique is the elbow method: run k-means for a range of k values, record the inertia (the sum of squared distances from each point to its assigned centroid) for each, and plot inertia against k. Inertia always decreases as k grows, but it typically drops sharply at first and then levels off; the 'elbow' where the curve bends is a reasonable choice for k.
Compared to hierarchical clustering, k-means scales much better to large datasets since it does not need a full pairwise distance matrix, but it forces you to commit to k upfront and assumes clusters are roughly round and similarly sized, an assumption hierarchical clustering does not need to make.
- Pros: fast and scalable to large datasets, simple to understand and implement.
- Cons: must choose k in advance, sensitive to centroid initialization, struggles with clusters of very different sizes, densities, or non-round shapes.
Exercise: Data Science Clustering
What type of machine learning technique is clustering?