K-Means partitions data into \(k\) clusters by iteratively placing centroids at the center of each group and reassigning points to their nearest centroid — repeating until the assignments stop changing.
The Objective Function
\(C_j\) is cluster \(j\)'s set of assigned points, and \(\mu_j\) is its centroid (mean). \(J\) — also called WCSS (Within-Cluster Sum of Squares) — sums the squared distance from every point to its own cluster's centroid. K-Means' entire algorithm is a search for centroid positions that minimize this quantity.
The Algorithm, As Explicit Steps
| Step | What Happens |
|---|---|
| 1 | Choose \(k\), and initialize \(k\) centroids (often randomly, or via smarter init like k-means++) |
| 2 | Assign: put each point in the cluster of its nearest centroid |
| 3 | Update: move each centroid to the mean of the points now assigned to it |
| 4 | Repeat steps 2-3 until assignments stop changing (convergence) |
Worked Example
6 points: \((2,2),(2,4),(2,6),(8,2),(8,4),(8,6)\), \(k=2\). Initialize centroids at \(\mu_1=(2,2)\), \(\mu_2=(8,6)\).
| Iteration | Assignment | New Centroids |
|---|---|---|
| 1 | Left 3 points → cluster 1, right 3 points → cluster 2 (each closer to its side) | \(\mu_1=(2,4)\), \(\mu_2=(8,4)\) |
| 2 | Same assignment (nothing changes) | Converged |
from sklearn.cluster import KMeans
import numpy as np
X = np.array([[2,2],[2,4],[2,6],[8,2],[8,4],[8,6]])
model = KMeans(n_clusters=2, init=np.array([[2,2],[8,6]]), n_init=1, random_state=42)
model.fit(X)
print(model.cluster_centers_) # [[2. 4.] [8. 4.]] -- matches the hand calculation
print(model.inertia_) # 16.0 -- scikit-learn's name for WCSS/J
print(model.labels_)
Graphical Intuition — Assignment and Centroid Movement
Centroids start at arbitrary points and drift toward the true center of their assigned group over iterations, until they stop moving.
Practical Use Cases
- Customer segmentation, market basket grouping, image color quantization
- A fast first-pass clustering method when clusters are expected to be roughly round/convex and similarly sized
Advantages
- Simple, fast, and scales well to large datasets
- Easy to interpret — each cluster has a clear, literal centroid representing its "average member"
Limitations
- Requires choosing \(k\) in advance — see the Elbow Method
- Assumes roughly round, similarly-sized clusters — struggles with irregular shapes (unlike DBSCAN)
- Sensitive to centroid initialization — different starting points can converge to different final clusterings, which is why
n_init(multiple random restarts) is used in practice - Sensitive to outliers, since the mean (used for centroids) is itself outlier-sensitive
Common Mistakes
- Running K-Means on unscaled features, letting one large-range feature dominate distance calculations.
- Using a single random initialization (
n_init=1) in production instead of multiple restarts, risking convergence to a poor local optimum.
Interview Relevance
Q: "Why does K-Means sometimes give different results across different runs on the same data?" Because it converges to a local optimum of a non-convex objective, and the local optimum reached depends on the random initial centroid positions — running multiple initializations (n_init) and keeping the best result (lowest WCSS) mitigates this.
Practice Question
Using the 6-point example, verify by hand that assigning \((8,2)\) to \(\mu_2=(8,4)\) rather than \(\mu_1=(2,4)\) is correct, by comparing the two squared distances.