K-Means Clustering Cheat Sheet
A reference for K-Means clustering covering scikit-learn implementation, centroid initialization, the elbow method, and silhouette scoring for choosing k.
Clustering with scikit-learn
Fit K-Means and inspect the resulting clusters.
from sklearn.cluster import KMeansfrom sklearn.preprocessing import StandardScalerX_scaled = StandardScaler().fit_transform(X)kmeans = KMeans(n_clusters=4, init='k-means++', n_init=10, random_state=42)labels = kmeans.fit_predict(X_scaled)print('Inertia:', kmeans.inertia_)print('Centroids:', kmeans.cluster_centers_)
Elbow Method
Pick k by plotting inertia across candidate values.
import matplotlib.pyplot as pltinertias = []for k in range(1, 11): km = KMeans(n_clusters=k, n_init=10, random_state=42).fit(X_scaled) inertias.append(km.inertia_)plt.plot(range(1, 11), inertias, marker='o')plt.xlabel('k'); plt.ylabel('Inertia') # look for the 'elbow' bend
Silhouette Score
Quantify cluster separation quality for each k.
from sklearn.metrics import silhouette_scorefor k in range(2, 8): labels = KMeans(n_clusters=k, n_init=10, random_state=42).fit_predict(X_scaled) score = silhouette_score(X_scaled, labels) print(f'k={k}: silhouette={score:.3f}') # closer to 1 is better
Key Concepts
Core theory behind K-Means.
- Centroid- Mean position of all points assigned to a cluster; recomputed every iteration
- Inertia- Sum of squared distances from points to their nearest centroid (within-cluster variance)
- k-means++- Smart centroid initialization that spreads out starting centroids to speed up convergence
- Elbow method- Plot inertia against k and pick the point where the decrease sharply flattens
- Silhouette score- Measures how similar a point is to its own cluster vs. neighboring clusters, ranging -1 to 1
- Convergence- Assignment and update steps alternate until centroids stop moving or max_iter is reached
MiniBatchKMeans for Large/Streaming Data
Cluster datasets too large for full-batch Lloyd's algorithm by updating centroids on random mini-batches.
from sklearn.cluster import MiniBatchKMeansmbk = MiniBatchKMeans( n_clusters=8, batch_size=1024, max_no_improvement=10, # early stop if inertia plateaus reassignment_ratio=0.01, # reseed rarely-used centroids random_state=42,)# Fit incrementally, e.g. from a generator of chunksfor chunk in stream_batches(X, size=1024): mbk.partial_fit(chunk)labels = mbk.predict(X)print('Inertia (approx):', mbk.inertia_)
Lloyd's Algorithm From Scratch
Implement the assign/update loop manually to see exactly what fit_predict is doing internally.
import numpy as npdef kmeans(X, k, max_iter=300, tol=1e-4, seed=42): rng = np.random.default_rng(seed) centroids = X[rng.choice(len(X), k, replace=False)] for _ in range(max_iter): # Assignment step: squared Euclidean distance to each centroid dists = ((X[:, None, :] - centroids[None, :, :]) ** 2).sum(axis=2) labels = dists.argmin(axis=1) # Update step: recompute centroids as cluster means new_centroids = np.array([ X[labels == j].mean(axis=0) if np.any(labels == j) else centroids[j] for j in range(k) ]) shift = np.linalg.norm(new_centroids - centroids) centroids = new_centroids if shift < tol: break return labels, centroids
Davies-Bouldin & Calinski-Harabasz Indices
Cross-check silhouette with label-only validation metrics that don't require pairwise distances.
from sklearn.metrics import davies_bouldin_score, calinski_harabasz_scorefor k in range(2, 8): labels = KMeans(n_clusters=k, n_init=10, random_state=42).fit_predict(X_scaled) db = davies_bouldin_score(X_scaled, labels) # lower is better, 0 is ideal ch = calinski_harabasz_score(X_scaled, labels) # higher is better print(f'k={k}: davies_bouldin={db:.3f} calinski_harabasz={ch:.1f}')
Dimensionality Reduction Before Clustering
Project onto principal components first so Euclidean distance stays meaningful in high dimensions, then tune the Lloyd variant.
from sklearn.decomposition import PCAfrom sklearn.pipeline import make_pipelinepipe = make_pipeline( StandardScaler(), PCA(n_components=0.95), # keep 95% of variance KMeans(n_clusters=5, algorithm='elkan', n_init=10, random_state=42),)labels = pipe.fit_predict(X)# 'elkan' exploits the triangle inequality to skip redundant distance# computations on dense, low-to-moderate dimensional data (default is 'lloyd')
Advanced Concepts
Beyond textbook K-Means: variants, limitations, and initialization theory.
- k-means|| (scalable k-means++)- Parallel initialization variant that samples multiple candidate centroids per round instead of one-at-a-time, used internally for large n_init runs
- algorithm='elkan' vs 'lloyd'- Elkan caches distance bounds via the triangle inequality to skip work; faster on dense low-dim data but more memory, ill-suited to sparse input
- Empty cluster handling- If a centroid captures zero points, scikit-learn reseeds it at the point farthest from its current centroid to avoid a degenerate solution
- Gap statistic- Compares within-cluster dispersion to that of a reference null (uniform) distribution to pick k more rigorously than the elbow heuristic
- Spherical/cosine k-means- Standard K-Means only supports Euclidean distance; for cosine similarity (e.g. text embeddings), normalize vectors to unit length first so Euclidean distance becomes rank-equivalent to cosine distance
- Kernel K-Means- Maps points into a higher-dimensional feature space via a kernel trick to separate non-convex clusters that vanilla K-Means cannot
- n_init='auto'- Modern scikit-learn defaults n_init to 1 run for k-means++ init (already good) vs 10 for random init, reducing redundant computation
K-means assumes roughly spherical, similarly sized clusters and is sensitive to feature scale and outliers — always standardize your features first, and consider DBSCAN or a Gaussian Mixture Model when clusters are non-convex or have very different densities.