K-Means clustering algorithm

What is K-Means?

WHY this objective? We want clusters where points are tightly packed around their center. Minimizing squared distances penalizes outliers heavily (quadratic penalty) and makes the math tractable for gradient-based reasoning.

The Algorithm: Derivation from First Principles

GOAL: Partition data $\{x_1, x_2, ..., x_n\}$ into $K$ clusters to minimize total variance.

Step 1: Why Centroids?

If we fix cluster assignments, what's the best representative point for cluster $C_k$?

Minimize: $\sum_{x_i \in C_k} ||x_i - \mu_k||^2$ with respect to $\mu_k$

Take the derivative and set to zero: $\frac{\partial}{\partial \mu_k} \sum_{x_i \in C_k} ||x_i - \mu_k||^2 = \sum_{x_i \in C_k} 2(x_i - \mu_k)(-1) = 0$

$\sum_{x_i \in C_k} (x_i - \mu_k) = 0$

$\sum_{x_i \in C_k} x_i = |C_k| \mu_k$

$\mu_k = \frac{1}{|C_k|} \sum_{x_i \in C_k} x_i$

RESULT: The optimal representative is the mean (average) of all points in the cluster. This is why we call them "centroids."

Step 2: Why Nearest-Neighbor Assignment?

If we fix centroids, which cluster should point $x_i$ join?

Minimize: The contribution of $x_i$ to the total objective $J$

Point $x_i$ should belong to cluster $k^*$ where: $k^* = \arg\min_{k \in \{1,...,K\}} ||x_i - \mu_k||^2$

WHY? Each point independently choses the centroid that minimizes its own squared distance. This gredy assignment is optimal when centroids are fixed.

Step 3: The Lloyd's Algorithm

These two steps are coupled: better assignments lead to better centroids, which lead to better assignments. We alternate:

Initialize: Randomly select $K$ points as initial centroids $\{\mu_1, ..., \mu_K\}$

Repeat until convergence:

1. Assignment Step: For each point $x_i$, assign to nearest centroid: $c_i = \arg\min_{k} ||x_i - \mu_k||^2$

2. Update Step: Recompute each centroid as the mean of assigned points: $\mu_k = \frac{1}{|C_k|} \sum_{x_i \in C_k} x_i$

Convergence: Stop when assignments no longer change (or centroids move less than threshold $\epsilon$)

WHY does this converge? Each step decreases (or keeps constant) the objective $J$:

• Assignment step: Points move to closer centroids → $J$ decreases
• Update step: Centroids move to the mean of their cluster → $J$ decreases (from our derivative proof)
• $J \geq 0$ is bounded below
• By monotone convergence theorem, algorithm must converge

IMPORTANT: Convergence is to a local minimum, not necessarily global. Different initializations can yield different results.

Worked Examples

Data: $A=(1,1), B=(2,1), C=(5,5), D=(6,6)$ Goal: Find 2 clusters

Iteration 0(Initialize):

• Random centroids: $\mu_1 = (1,1), \mu_2 = (5,5)$

**Iteration 1:*Assignment:

• $A=(1,1)$: $d(\mu_1)=0, d(\mu_2)=\sqrt{32}$ → Cluster 1
• $B=(2,1)$: $d(\mu_1)=1, d(\mu_2)=\sqrt{32}$ → Cluster 1
• $C=(5,5)$: $d(\mu_1)=\sqrt{32}, d(\mu_2)=0$ → Cluster 2
• $D=(6,6)$: $d(\mu_1)=\sqrt{50}, d(\mu_2)=\sqrt{2}$ → Cluster 2

Update:

• $\mu_1 = \frac{(1,1)+(2,1)}{2} = (1.5, 1)$
• $\mu_2 = \frac{(5,5)+(6,6)}{2} = (5.5, 5.5)$

Iteration 2:

Assignment: (Check distances again)

• All assignments remain the same

Converged! Final clusters: $\{A, B\}$ and $\{C, D\}$

WHY THIS STEP? Each distance calculation uses Euclidean metric because we're minimizing squared distances. The update step averages coordinates component-wise because the mean minimizes sum of squared deviations.

Problem: Cluster 100 customers by [spending, frequency]

# Pseudocode with reasoning
data = load_customers()  # Shape: (100, 2)
K = 3
 
# Initialize: K-Means++ (smarter than random)
centroids = [random_point()]
for in range(K-1):
    # Pick next centroid far from existing ones
    distances = [min(dist(p, c) for c in centroids) for p in data]
    centroids.append(weighted_random(data, distances))
 
for iteration in range(max_iters):
    # Assignment: Vectorized for speed
    distances = cdist(data, centroids)  # Shape (100, 3)
    labels = argmin(distances, axis=1)   # Shape (100,)
    # Update
    new_centroids = []
    for k in range(K):
        cluster_points = data[labels == k]
        new_centroids.append(mean(cluster_points, axis=0))
    
    if converged(centroids, new_centroids):
        break
    centroids = new_centroids

WHY THESE CHOICES?

• K-Means++ initialization: Spreads initial centroids out → faster convergence, avoids bad local minima
• Vectorization: cdist computes all pairwise distances at once → 100× faster than loops
• Convergence check: Compare centroid movement, not assignments → more stable stopping criterion

Common Mistakes & How to Fix Them

The reality: K-Means converges to a local minimum, which depends on initialization. Different random starts give different results.

Steel-man: The algorithm does optimize the WCSS objective perfectly given the starting point. The issue is the objective landscape has many valleys.

Fix:

• Run K-Means multiple times with different initializations (sklearn default: 10 runs)
• Use K-Means++ initialization (provably better than random)
• Report the solution with lowest final WCSS

The reality: K-Means assumes spherical clusters of similar size. It fails on:

• Elongated/elliptical clusters
• Crescent or ring shapes
• Clusters of very different densities

Steel-man: For data generated from spherical Gaussian distributions, K-Means is actually optimal (it's the MLE solution).

Fix:

• Preprocess: Standardize features if scales differ
• Use DBSCAN or spectral clustering for non-convex shapes
• Transform data (e.g., kernel trick) to make clusters more spherical

The reality:

• WCSS always decreases as $K$ increases (at $K=n$, WCSS=0)
• No universal "correct" K—it's domain-dependent

Steel-man: Systematic search is valuable, but you need the right metric.

Fix:

• Elbow method: Plot WCSS vs K, look for "elbow" where improvement slows
• Silhouette score: Measures how similar points are to their own cluster vs. others
• Domain knowledge: Sometimes K is known (e.g., customer segments)
• Gap statistic: Compare WCSS to null distribution

Computational Complexity

Time per iteration:

• Assignment step: $O(nKd)$ — compute $n$ points × $K$ centroids × $d$ dimensions
• Update step: $O(nd)$ — sum over $n$ points, $d$ dimensions

Total: $O(IKnd)$ where $I$ = number of iterations

Space: $O((n+K)d)$ — store data and centroids

In practice: $I$ is typically small (<100) and $K \ll n$, so K-Means scales well to large datasets.

WHY this complexity? The assignment step is the bottleneck: we compute $nK$ distances. Advanced implementations use KD-trees or ball trees to reduce this when $d$ is small.

The Math Behind Convergence

Claim: Each iteration decreases $J$ (or keeps it the same), so K-Means converges.

Proof: Let $J^{(t)}$ be the objective at iteration $t$.

After assignment step at iteration $t+1$: $J_{\text{assign}}^{(t+1)} \leq J^{(t)}$ because each point moves to the nearest centroid (or stays), which cannot increase its distance.

After update step: $J^{(t+1)} \leq J_{\text{assign}}^{(t+1)}$ because $\mu_k = \arg\min_mu \sum_{x_i \in C_k} ||x_i - \mu||^2$ (proved earlier).

Therefore: $J^{(t+1)} \leq J^{(t)}$

Since $J \geq 0$ and there are finitely many possible labelings (at most $K^n$), the algorithm must converge. $\square$

CAVEAT: Convergence is to a local minimum. The global minimum is NP-hard to find.

Extensions and Variations

1. K-Means++: Smart initialization that spreads initial centroids $P(\text{choose } x_i) \propto D(x_i)^2$ where $D(x_i)$ = distance to nearest existing centroid

2. Mini-Batch K-Means: Use random subsets of data each iteration → faster for huge datasets

3. Soft K-Means (Fuzzy C-Means): Allow partial membership: $p(C_k | x_i) = \frac{e^{-\beta ||x_i - \mu_k||^2}}{\sum_j e^{-\beta ||x_i - \mu_j||^2}}$

4. K-Medoids (PAM): Use actual data points as centers (robust to outliers)

1. You put down K different boxes randomly in the pile
2. Each candy walks to the nearest box
3. You move each box to the center of all the candies that chose it
4. Repeat steps 2-3 until the boxes stop moving

The boxes are "learning" where the groups of similar candies naturally are! The candies in the same box end up being similar to each other, even though you never told the algorithm which candies should go together.

The tricky part is you have to decide how many boxes (K) to use at the start. Too few boxes and you squish different types together. Too many boxes and you split up candies that are actually the same type.

"Converge but Local" — It stops moving, but might not be the best answer overall

K-Means = K-Averages — Centroids are literally means (averages), not medians or modes

"Spherical, Similar Size" — The two S's that describe what K-Means likes

Connections

• Choosing K: Elbow & Silhouette Methods
• K-Means++ Initialization
• Hierarchical Clustering — alternative approach without predefining K
• DBSCAN — handles non-spherical clusters
• Gradient Descent — similar alternating optimization pattern
• Bias-Variance Tradeoff — choosing K trades off fit vs. generalization
• Expectation-Maximization (EM) — K-Means is a special case (hard EM for Gaussians)

#flashcards/ai-ml

What is the objective function that K-Means minimizes? :: Within-cluster sum of squares (WCSS): $J = \sum_{k=1}^{K} \sum_{x_i \in C_k} ||x_i - \mu_k||^2$

Concept Map

Hinglish (regional understanding)