# K-Means Clustering

Partitioning data into k groups by alternately assigning points to the nearest centroid and recomputing centroids.

**Rests on:** Dimensionality Reduction and PCA

*Machine Learning · Unsupervised Learning*


## Why learn this

K-means clustering lets you find groups in data without anyone telling you what the groups are. It is the first algorithm you meet in unsupervised learning — the branch of machine learning where there is no 'correct answer' to check against. Master this, and you unlock the rest of unsupervised learning: hierarchical clustering (where you build a tree of groups instead of picking a number in advance), DBSCAN and OPTICS (which find oddly-shaped clusters and label stray points as noise), and PCA (which compresses many features into a few). On top of that, k-means is the gateway to embeddings — the vectors that large language models use to represent words, sentences and documents. When you retrieve a passage from a knowledge base to answer a question in a RAG system, embeddings are often clustered with something like k-means to organise them. When you hear about 'semantic search' or 'similarity', that machinery often rests on ideas you will meet here.


## The intuition

K-means answers: given a pile of points, divide them into k groups so that points in the same group are close together, and points in different groups are far apart. Imagine you are a teacher sorting students into k study groups by ability. You might start by picking k students as 'representatives' (centroids). Then each student joins the nearest representative's group. After that, you compute the average of each group and move the representative there. You repeat: reassign every student to the nearest representative, recompute the averages, until no student changes group. That is k-means. It works when groups are roughly spherical (like blobs) and roughly equal in size. If a group is crescent-shaped, k-means will cut it in half. If one group is ten times bigger than another, k-means may split the big one and merge the small ones. It is simple, fast, and the go-to baseline for clustering — but its assumptions are often false in real data. Your library notes emphasise this: 'assumes spherical, similar-sized clusters — often false and rarely checked'. K-means is not a magic wand; it is a quick tool that works well when the blobs are blobs.


## An analogy

Think of a department store sorting its customers into k loyalty segments. The store knows each customer's average monthly spend and how many visits they make. They want to divide customers into k groups for targeted marketing. First, they pick k employees to be 'representative customers' — maybe one high-spender, one low-spender, one medium. Each customer is assigned to the employee most similar in spend and visits. The store then computes the average spend and visits for each group and promotes the 'average customer' of each group to be the new representative. They repeat: reassign customers to the nearest representative (now an average), compute new averages, until no one switches. The store ends up with k groups: 'bargain hunters', 'loyalists', 'occasional splurgers', etc. The analogy works because k-means is exactly this: assign to nearest, recompute prototypes, repeat. Where it breaks: real customer segments are not necessarily spherical blobs in spend-visit space. A 'bargain hunter' might visit often but spend little, or rarely but stock up on discounts — those form different shapes. K-means would force them into one or two blobs, possibly mixing distinct behaviours. Also, the store chose k in advance — but how many segments really exist? The elbow method (which we'll see) is a guess, not an answer. And the starting representatives matter: choose differently, get a different clustering — k-means++ reduces this risk but doesn't eliminate it.


## Definition

K-means clustering partitions a set of points into k groups by (1) assigning each point to the nearest of k initial centroids, (2) recomputing each centroid as the average of its assigned points, and repeating until assignments stop changing.


## How it connects

You have not met any machine learning yet — this is your first. That is fine; k-means is the gentlest entry point. It belongs to unsupervised learning, because there is no target variable to predict; you are finding structure in unlabelled data. Your library lists hierarchical clustering, DBSCAN/OPTICS, and PCA as neighbours. Hierarchical clustering builds a tree of nested clusters, so you don't need to pick k in advance — you cut the tree afterwards. DBSCAN drops k-means' spherical assumption and labels noise points. PCA compresses many features into a few, which is often done before clustering to speed it up and reduce noise. K-means also connects to embeddings later: when you embed text into vectors, you might cluster those vectors to organise a knowledge base. So this is not an isolated trick; it is a foundation stone.


## The mathematics

K-means clustering is a method for grouping data points into clusters. It works by repeatedly adjusting the center of each group until the groups stop changing. The math it uses is simple arithmetic: adding numbers, dividing them to find averages, and calculating distances between points.

The core idea is to minimize the total distance from every point to the center of its own cluster. To do this, the algorithm alternates between two steps:

1. Assignment Step: For each data point, calculate how far it is from each cluster center. Assign the point to the cluster whose center is closest.
2. Update Step: For each cluster, calculate the new center. The new center is the average (mean) of all the points currently assigned to that cluster.

We repeat these two steps until no point changes its cluster assignment. At that point, the algorithm has converged.

Distance Calculation
The standard way to measure distance in K-means is Euclidean distance. This is the straight-line distance between two points. For data with more than two dimensions (for example, if each point has an x-coordinate, a y-coordinate, and a z-coordinate), we use the generalized Pythagorean theorem.

Suppose a data point has coordinates (x1, x2, ..., xd) and a cluster center has coordinates (c1, c2, ..., cd). Here, d is the number of dimensions. The squared Euclidean distance between them is calculated by:
- Subtracting each coordinate of the center from the corresponding coordinate of the point: (x1 - c1), (x2 - c2), etc.
- Squaring each of these differences: (x1 - c1)^2, (x2 - c2)^2, etc.
- Adding all the squared differences together.

The formula for the squared distance is:
distance_squared = (x1 - c1)^2 + (x2 - c2)^2 + ... + (xd - cd)^2

We use squared distance because it avoids the need to take square roots during the assignment step. Since the square root function increases steadily, the point with the smallest squared distance is also the point with the smallest actual distance. This makes the calculation faster without changing which cluster a point belongs to.

Updating the Center
Once points are assigned to clusters, we update the center of each cluster. The new center is the mean (average) of all points in that cluster. For each dimension separately:
- Add up all the values for that dimension across all points in the cluster.
- Divide by the number of points in the cluster.

For example, if a cluster has three points with x-coordinates 1, 2, and 5, the new x-coordinate for the center is (1 + 2 + 5) / 3 = 8 / 3 ≈ 2.67.

Why this works
Each step reduces the total within-cluster variance. The assignment step ensures every point is as close as possible to its current center. The update step ensures the center is exactly in the middle of its points, minimizing the sum of squared distances for that cluster. Because the total distance can only decrease or stay the same at each step, and it cannot go below zero, the process must eventually stop changing.

Note on Robustness
K-means uses the mean, which is sensitive to outliers. A single point far away from the rest of the cluster will pull the center toward it. This is a known limitation of the method.


## A worked example

Let us walk through one full iteration of K-means clustering with six points in two dimensions.

The Data Points:
- Point A: (1, 2)
- Point B: (2, 1)
- Point C: (5, 4)
- Point D: (6, 5)
- Point E: (8, 8)
- Point F: (9, 9)

We want to find k = 2 clusters. We start by choosing two initial centers (centroids). Let us pick:
- Center C1: (1, 2) [This is the same location as Point A]
- Center C2: (8, 8) [This is the same location as Point E]

Step 1: Assignment Phase
We calculate the squared Euclidean distance from each point to both centers. We use squared distance because it is easier to compute and gives the same result for finding the minimum.

Formula for squared distance between (x, y) and (cx, cy):
dist_sq = (x - cx)^2 + (y - cy)^2

Point A (1, 2):
- Distance to C1 (1, 2): (1-1)^2 + (2-2)^2 = 0 + 0 = 0
- Distance to C2 (8, 8): (1-8)^2 + (2-8)^2 = (-7)^2 + (-6)^2 = 49 + 36 = 85
- Closest center: C1 (distance 0 < 85). Assign A to Cluster 1.

Point B (2, 1):
- Distance to C1 (1, 2): (2-1)^2 + (1-2)^2 = 1^2 + (-1)^2 = 1 + 1 = 2
- Distance to C2 (8, 8): (2-8)^2 + (1-8)^2 = (-6)^2 + (-7)^2 = 36 + 49 = 85
- Closest center: C1 (distance 2 < 85). Assign B to Cluster 1.

Point C (5, 4):
- Distance to C1 (1, 2): (5-1)^2 + (4-2)^2 = 4^2 + 2^2 = 16 + 4 = 20
- Distance to C2 (8, 8): (5-8)^2 + (4-8)^2 = (-3)^2 + (-4)^2 = 9 + 16 = 25
- Closest center: C1 (distance 20 < 25). Assign C to Cluster 1.

Point D (6, 5):
- Distance to C1 (1, 2): (6-1)^2 + (5-2)^2 = 5^2 + 3^2 = 25 + 9 = 34
- Distance to C2 (8, 8): (6-8)^2 + (5-8)^2 = (-2)^2 + (-3)^2 = 4 + 9 = 13
- Closest center: C2 (distance 13 < 34). Assign D to Cluster 2.

Point E (8, 8):
- Distance to C1 (1, 2): (8-1)^2 + (8-2)^2 = 7^2 + 6^2 = 49 + 36 = 85
- Distance to C2 (8, 8): (8-8)^2 + (8-8)^2 = 0 + 0 = 0
- Closest center: C2 (distance 0 < 85). Assign E to Cluster 2.

Point F (9, 9):
- Distance to C1 (1, 2): (9-1)^2 + (9-2)^2 = 8^2 + 7^2 = 64 + 49 = 113
- Distance to C2 (8, 8): (9-8)^2 + (9-8)^2 = 1^2 + 1^2 = 1 + 1 = 2
- Closest center: C2 (distance 2 < 113). Assign F to Cluster 2.

After this assignment step, the clusters are:
- Cluster 1 contains points: A, B, C
- Cluster 2 contains points: D, E, F

Step 2: Update Phase
Now we recalculate the center for each cluster by taking the mean of the coordinates of the points in that cluster.

New Center for Cluster 1 (C1_new):
Points are A(1,2), B(2,1), C(5,4).
- New x-coordinate: (1 + 2 + 5) / 3 = 8 / 3 ≈ 2.67
- New y-coordinate: (2 + 1 + 4) / 3 = 7 / 3 ≈ 2.33
So, C1_new is approximately (2.67, 2.33).

New Center for Cluster 2 (C2_new):
Points are D(6,5), E(8,8), F(9,9).
- New x-coordinate: (6 + 8 + 9) / 3 = 23 / 3 ≈ 7.67
- New y-coordinate: (5 + 8 + 9) / 3 = 22 / 3 ≈ 7.33
So, C2_new is approximately (7.67, 7.33).

End of Iteration 1:
The centers have moved from (1,2) and (8,8) to (2.67, 2.33) and (7.67, 7.33). In a real run, we would go back to Step 1 and recalculate distances using these new centers. Some points might switch clusters if they are now closer to the other center. We repeat this process until no points change their cluster assignment.


## From scratch, in Python

```python
# Plain Python implementation of k-means
# We'll use lists of points, each point is a tuple (x, y)

# Our data: six points as a list of tuples
points = [(1,2), (2,1), (5,4), (6,5), (8,8), (9,9)]

# Choose k=2 and initial centroids (we'll pick the first and fifth points)
k = 2
centroids = [points[0], points[4]]  # (1,2) and (8,8)

# A helper to compute Euclidean distance between two points
def distance(p, q):
    # p and q are tuples like (x,y)
    # square the differences, sum, then square root
    return ((p[0]-q[0])**2 + (p[1]-q[1])**2) ** 0.5

# Main loop: repeat until no point changes cluster
while True:
    # Step 1: assign each point to nearest centroid
    clusters = [[] for _ in range(k)]  # list of empty lists, one per cluster
    for point in points:
        # find index of nearest centroid
        nearest = 0
        best_dist = distance(point, centroids[0])
        for j in range(1, k):
            d = distance(point, centroids[j])
            if d < best_dist:
                best_dist = d
                nearest = j
        clusters[nearest].append(point)
    
    # Step 2: recompute centroids as averages
    new_centroids = []
    for cluster in clusters:
        # compute average x and average y
        n = len(cluster)
        if n == 0:
            # avoid division by zero; keep old centroid (not in our example)
            new_centroids.append(centroids[clusters.index(cluster)])
        else:
            sum_x = sum(p[0] for p in cluster)
            sum_y = sum(p[1] for p in cluster)
            new_centroids.append((sum_x/n, sum_y/n))
    
    # Check if centroids changed (compare with some small tolerance)
    changed = False
    for old, new in zip(centroids, new_centroids):
        if distance(old, new) > 1e-9:
            changed = True
            break
    if not changed:
        break
    centroids = new_centroids

print("Final clusters:")
for i, cluster in enumerate(clusters):
    print(f"Cluster {i}: {cluster}")
print("Final centroids:")
for c in centroids:
    print(c)
```


## The same thing, with a framework

```python
import numpy as np
from sklearn.cluster import KMeans

# Points as a 2D numpy array
X = np.array([(1,2), (2,1), (5,4), (6,5), (8,8), (9,9)])

# The standard library call
kmeans = KMeans(n_clusters=2, init='k-means++', n_init=10, random_state=0)
kmeans.fit(X)

print(kmeans.labels_)           # which cluster each point belongs to
print(kmeans.cluster_centers_)  # the final centroids
```


## What the code does

The plain Python code is the heart. We first define our points as a list of tuples — a tuple is like a list but cannot be changed; it's a fixed pair (x, y). We choose k=2 and set initial centroids to the first and fifth points (just as an example). The `distance` function uses the Pythagorean formula: it squares the difference in x, squares the difference in y, adds, and takes the square root. The main loop runs until no centroid changes. Inside, we create `clusters` as a list of empty lists (one for each cluster). For each point, we loop over centroids to find the nearest (smallest distance). We append the point to that cluster's list. Then, for each cluster, we compute the average x and average y by summing all x's and dividing by the count, and similarly for y. That gives the new centroid. We compare old and new centroids; if any moved more than a tiny tolerance, we keep looping. When nothing moves, we print the clusters and centroids. The framework version uses `sklearn.cluster.KMeans`, which does exactly the same assignment and update steps underneath, but with better initialisation (k-means++) and multiple restarts to avoid poor local optima. `kmeans.labels_` gives the cluster index for each input point, and `kmeans.cluster_centers_` gives the final centroid coordinates.


## Common misconceptions

**K-means finds the 'true' clusters.** — If I run k-means, it will discover the natural groups in my data.

K-means always produces k clusters, even if the data has no clusters or a different number. The clusters are spherical and similar-sized, so if the true groups are elongated or uneven, k-means will force a wrong partition. In fact: K-means finds a good k-way partition of the data under its assumptions. There is no ground truth in clustering, so the result is only one possible grouping. You must assess whether it makes sense for your problem.

**The elbow method picks the best k.** — The elbow method tells me exactly how many clusters to use.

The elbow method is a heuristic: you look for an 'elbow' in the within-cluster sum of squares plot, but the elbow is often ambiguous or absent. The number of clusters is a modelling choice, not a discovered truth. In fact: The elbow method gives a suggestion, but you should also consider domain knowledge, interpretability, and validation metrics like silhouette score. There is no automatic 'true k'.

**Initialisation doesn't matter.** — K-means will find the same clusters no matter how I start.

K-means uses random initial centroids, and the algorithm converges to a local optimum. Different starting points can give different final clusters. Poor initialisation can produce bad clusters. In fact: Use k-means++ initialisation (which spreads out initial centroids) and run multiple restarts (n_init). This makes results far more stable, but still not guaranteed global optimum.

**K-means works on any data shape.** — If I have two crescents, k-means will separate them nicely.

K-means assigns each point to the nearest centroid, dividing space by straight lines (Voronoi cells). A crescent shape cannot be captured this way; k-means will cut across the crescent, producing wrong clusters. In fact: K-means assumes convex, roughly spherical clusters. If your data has arbitrary shapes, use density-based methods like DBSCAN or OPTICS, which find arbitrarily shaped clusters and label noise.

**The number of clusters k is given by the data.** — I can just run k-means with k=3 and that's the right answer.

K-means does not tell you what k should be. k is a hyperparameter you choose. The algorithm will always partition the data into exactly k groups, regardless of what is natural. In fact: You must choose k based on domain knowledge, heuristic methods (elbow, silhouette), and the downstream task. There is no universal answer.

**K-means is only for 2D data.** — K-means can only handle 2D points.

K-means works in any number of dimensions. The Euclidean distance formula generalises to D dimensions. High-dimensional data is common (e.g., embeddings of 1000 dimensions). In fact: K-means works in any D-dimensional space; the algorithm is the same. However, the curse of dimensionality can make distances less meaningful in very high dimensions, so dimensionality reduction (PCA) is often used first.

**K-means is slow and impractical.** — For millions of data points, k-means would take forever.

K-means is one of the fastest clustering algorithms. Its time complexity is roughly O(n * k * I * D), where n is number of points, k clusters, I iterations, D dimensions. For n=1 million, k=100, I=100, D=100, that's about 10^12 operations, but optimised implementations with vectorisation can handle this in seconds to minutes. In fact: K-means is designed to be scalable. Libraries like scikit-learn, FAISS, and Spark implement efficient versions. On GPUs, it can be even faster.


## In short

K-means clustering groups data into k clusters by alternately assigning points to the nearest centroid and recomputing centroids as averages. It is simple, fast, and widely used, but it assumes spherical, similar-sized clusters, requires choosing k in advance, and depends on initialisation. Use k-means++ and multiple restarts for stability. It is a foundational tool in unsupervised learning, with applications in customer segmentation, image compression, vector search (e.g., RAG), and model compression.


## Where this is used

K-means is a workhorse in applied machine learning and appears in many AI systems, often behind the scenes. Here is how, in detail.

**Large language models (LLMs).** LLMs represent words, sentences, and documents as vectors (embeddings) in a high-dimensional space. K-means is sometimes used to cluster these embeddings for efficient retrieval. For instance, in retrieval-augmented generation (RAG), a knowledge base is chunked into passages, each embedded into a vector. Instead of doing a brute-force search over every passage for each query (which is slow), you can cluster the passages with k-means. When a query arrives, you first find which cluster(s) the query's embedding is closest to, and then search only within those. This is called 'cluster-based retrieval' or 'inverted file index'. The mechanism: you train a k-means model on the passage embeddings, get cluster centers, and map each passage to its cluster ID. At query time, you embed the query, compute distances to all cluster centers (a small number, e.g., 100), pick the nearest few clusters, and then compute similarity to only the passages in those clusters. This reduces the search space dramatically. A developer would write something like: `kmeans = KMeans(n_clusters=100); kmeans.fit(all_embeddings);` and then store `kmeans.labels_` alongside each passage.

**Small models on devices.** When you have a small language model running on a phone or a smart speaker (e.g., for keyword spotting or on-device translation), k-means is used for vector quantization of model weights or activations. The idea: the weights are floating-point numbers that take up space. K-means can find, say, 256 representative values (centroids) and replace each weight with the index of its nearest centroid. This compresses the model (e.g., from 32-bit floats to 8-bit indices) with little loss in accuracy, because many weights are similar. This is called 'k-means quantization'. The developer would run k-means on the weight values, then store the codebook (centroids) and the indices. On the device, the model looks up the centroid values instead of storing the full numbers, reducing memory and speeding up computation.

**Retrieval-augmented generation (RAG).** Beyond just clustering passages for speed, k-means is also used to organise the embedding space for better retrieval. Sometimes, you group related documents into clusters and then for each cluster you train a smaller model or a summary, so that when a query hits a cluster, you can quickly retrieve a summary or a sub-index. In practice, developers use k-means from scikit-learn on the embedding vectors, often after reducing dimensionality with PCA (your library mentions PCA as a neighbour). A common pattern: `from sklearn.cluster import KMeans; km = KMeans(n_clusters=params['k']); km.fit(doc_embeddings);` then build an inverted index mapping each cluster to its members. When building the system, you must choose k — often via the elbow method, but that is heuristic. The clustering quality directly affects retrieval accuracy; if clusters are bad, you might miss relevant passages. So developers often tune k and use multiple restarts (n_init) to avoid poor local optima.

**Agentic systems.** Agents that plan and call tools often need to decide which tool to use or which memory to retrieve. K-means can cluster the agent's memory (past experiences, instructions, etc.) so that when a new situation arises, the agent can quickly find the most relevant cluster. The mechanism: the agent maintains a memory of embedding vectors of past steps. Clustering these memories allows the agent to look up similar past situations by first finding the cluster nearest to the current state, then searching within that cluster. This speeds up retrieval and can also be used to summarise experiences per cluster, giving the agent a 'memory of clusters' rather than raw memories. In a coding agent, for example, code snippets might be clustered by functionality; when a new task appears, the agent retrieves cluster centroids to decide which cluster of code to examine. The developer might use k-means from a library like FAISS (which supports k-means for large-scale clustering).

**Everyday applied ML.** K-means is used for customer segmentation (as in our analogy), image compression (reducing colours to k representative colours), anomaly detection (points far from any centroid are suspicious), and pre-processing for other algorithms. In marketing, you cluster customers by purchase history and demographics; in vision, you cluster image patches into visual words for bag-of-words models. A developer working on a recommendation system might cluster users or items to generate recommendations; the mechanism is to compute user or item embeddings (e.g., from ratings) and then cluster, and for a new user, find the nearest cluster and recommend items popular in that cluster. In fraud detection, k-means can group transactions into normal and suspicious clusters; outliers (points far from centroids) are flagged. The standard library is scikit-learn in Python; for large-scale, FAISS or Spark's MLlib. What breaks when done badly: ignoring the spherical assumption leads to poor clusters; choosing k wrongly gives misleading groups; poor initialisation gives unstable results — so always set `n_init` to at least 10, use k-means++, and validate with silhouette scores if possible.

**Limitations and honest absence.** K-means is not used in training the weights of neural networks (that is gradient descent). It is not the core of attention mechanisms. It appears as a tool, not as a fundamental component of deep learning. But in any system that deals with large collections of vectors — which includes modern RAG and agent memory systems — clustering is a standard engineering technique. So it deserves its place as a foundational algorithm.


---

Read this on the web: https://basavapurushottam.com/learn-ai/k-means-clustering
