← Back

K-Means Clustering: 2D Vector Example

8 min read

In vector databases and search engines like Turbopuffer, K-Means clustering is the mathematical engine used to partition high-dimensional vector spaces into localized spatial clusters (often referred to as an Inverted File Index or IVF). By grouping similar vectors together, the database can restrict its search query to only the most relevant clusters, avoiding an expensive exhaustive search over the entire database.

This document walks through the mathematical formulation of K-Means clustering, performs a step-by-step hand-calculation using a simple 2D dataset, and provides clean code implementations in Python and TypeScript.

1. Mathematical Formula

The objective of K-Means clustering is to partition NN observations (vectors) x1,x2,…,xN∈RD\mathbf{x}_1, \mathbf{x}_2, \dots, \mathbf{x}_N \in \mathbb{R}^D into KK disjoint clusters S={S1,S2,…,SK}\mathbf{S} = \{S_1, S_2, \dots, S_K\} so as to minimize the Within-Cluster Sum of Squares (WCSS), also known as Inertia:

J=∑j=1K∑xi∈Sj∥xi−cj∥2J = \sum_{j=1}^{K} \sum_{\mathbf{x}_i \in S_j} \|\mathbf{x}_i - \mathbf{c}_j\|^2

where cj\mathbf{c}_j is the mean (centroid) of the points in cluster SjS_j:

cj=1∣Sj∣∑xi∈Sjxi\mathbf{c}_j = \frac{1}{|S_j|} \sum_{\mathbf{x}_i \in S_j} \mathbf{x}_i

The algorithm achieves this by alternating between two main steps until the centroids converge (stop moving).

Step 1: Assignment Step (Voronoi Partitioning)

Each data point xi\mathbf{x}_i is assigned to the cluster with the nearest centroid. This step geometrically partitions the space into Voronoi cells:

Sj(t)={xi:∥xi−cj(t)∥2≤∥xi−cl(t)∥2 ∀ l, 1≤l≤K}S_j^{(t)} = \left\{ \mathbf{x}_i : \left\| \mathbf{x}_i - \mathbf{c}_j^{(t)} \right\|^2 \le \left\| \mathbf{x}_i - \mathbf{c}_l^{(t)} \right\|^2 \ \forall\ l, \ 1 \le l \le K \right\}

Step 2: Update Step (Centroid Re-calculation)

The coordinates of each centroid are updated to be the arithmetic mean of all data points assigned to that cluster in the assignment step:

cj(t+1)=1∣Sj(t)∣∑xi∈Sj(t)xi\mathbf{c}_j^{(t+1)} = \frac{1}{\left|S_j^{(t)}\right|} \sum_{\mathbf{x}_i \in S_j^{(t)}} \mathbf{x}_i

2. Walkthrough: 2D Numerical Example

Let's trace K-Means clustering manually using a simple two-dimensional dataset (D=2D = 2) and K=2K = 2 clusters.

The Dataset

We have N=5N = 5 points on a 2D coordinate grid:

We wish to group these points into K=2K = 2 clusters.

Iteration 1

Step 1.1: Initialization

We select two initial starting points to act as our centroids. Notice that both start on the right side of the grid:

Step 1.2: Distance Calculation & Assignment

We compute the squared Euclidean distance from each point to both centroids:

d2(p,c)=(xp−xc)2+(yp−yc)2d^2(\mathbf{p}, \mathbf{c}) = (x_p - x_c)^2 + (y_p - y_c)^2

Our partitioned clusters at t=1t=1 are:

S1(1)={C},S2(1)={A,B,D,E}S_1^{(1)} = \{\mathbf{C}\}, \quad S_2^{(1)} = \{\mathbf{A}, \mathbf{B}, \mathbf{D}, \mathbf{E}\}

Step 1.3: Centroid Update

We recalculate the centroids as the mean of their assigned points:

Iteration 2

We repeat the assignment step using our updated centroids:

Step 2.1: Distance Calculation & Re-assignment

Our partitioned clusters at t=2t=2 are:

S1(2)={C,D},S2(2)={A,B,E}S_1^{(2)} = \{\mathbf{C}, \mathbf{D}\}, \quad S_2^{(2)} = \{\mathbf{A}, \mathbf{B}, \mathbf{E}\}

Step 2.2: Centroid Update

Iteration 3

Step 3.1: Distance Calculation & Re-assignment

Our partitioned clusters at t=3t=3 are:

S1(3)={B,C,D},S2(3)={A,E}S_1^{(3)} = \{\mathbf{B}, \mathbf{C}, \mathbf{D}\}, \quad S_2^{(3)} = \{\mathbf{A}, \mathbf{E}\}

Step 3.2: Centroid Update

Iteration 4

Step 4.1: Distance Calculation & Re-assignment

Our partitioned clusters at t=4t=4 are:

S1(4)={B,C,D},S2(4)={A,E}S_1^{(4)} = \{\mathbf{B}, \mathbf{C}, \mathbf{D}\}, \quad S_2^{(4)} = \{\mathbf{A}, \mathbf{E}\}

Since the point assignments have not changed from the previous iteration, recalculating the centroids will result in the exact same coordinates:

The algorithm has converged after 4 iterations! The final clusters are the right-hand column {B,C,D}\{\mathbf{B}, \mathbf{C}, \mathbf{D}\} and the bottom-left pair {A,E}\{\mathbf{A}, \mathbf{E}\} — even though both centroids started on the right side of the grid.

3. Implementations

Here are clean, reference implementations of the K-Means algorithm for 2D vectors.

Python (NumPy)

python
import numpy as np
 
def kmeans_2d(points, k, max_iters=100, tol=1e-4):
    # Convert input to numpy array
    X = np.array(points, dtype=np.float64)
 
    # 1. Initialize centroids randomly from the dataset points
    rng = np.random.default_rng()
    initial_idx = rng.choice(X.shape[0], size=k, replace=False)
    centroids = X[initial_idx]
 
    for iteration in range(max_iters):
        # 2. Assignment Step: compute pairwise Euclidean distances
        # X[:, np.newaxis, :] is shape (N, 1, 2)
        # centroids[np.newaxis, :, :] is shape (1, K, 2)
        distances = np.linalg.norm(X[:, np.newaxis, :] - centroids[np.newaxis, :, :], axis=2)
        labels = np.argmin(distances, axis=1)
 
        # 3. Update Step: compute new centroids
        new_centroids = np.array([
            X[labels == j].mean(axis=0) if np.any(labels == j) else centroids[j]
            for j in range(k)
        ])
 
        # Check for convergence (if centroids move less than tolerance)
        if np.all(np.linalg.norm(new_centroids - centroids, axis=1) < tol):
            print(f"Converged at iteration {iteration}")
            break
 
        centroids = new_centroids
 
    return centroids, labels
 
# Test the example
data_points = [[2, 1], [6, 5], [6, 1], [6, 3], [1, 1]]
centroids, labels = kmeans_2d(data_points, k=2)
print("Centroids:\n", centroids)
print("Labels:\n", labels)

TypeScript / JavaScript

typescript
interface Point2D {
  x: number;
  y: number;
}
 
function calculateDistanceSquared(p1: Point2D, p2: Point2D): number {
  return Math.pow(p1.x - p2.x, 2) + Math.pow(p1.y - p2.y, 2);
}
 
function kmeans2D(
  points: Point2D[],
  k: number,
  maxIters = 100,
): { centroids: Point2D[]; labels: number[] } {
  // 1. Initialize centroids: take first K unique points (for simple deterministic seed)
  const centroids: Point2D[] = points.slice(0, k).map((p) => ({ ...p }));
  const labels: number[] = new Array(points.length).fill(-1);
 
  for (let iter = 0; iter < maxIters; iter++) {
    let changed = false;
 
    // 2. Assignment Step
    for (let i = 0; i < points.length; i++) {
      const p = points[i];
      let minDistance = Infinity;
      let closestCentroidIndex = -1;
 
      for (let j = 0; j < k; j++) {
        const distSq = calculateDistanceSquared(p, centroids[j]);
        if (distSq < minDistance) {
          minDistance = distSq;
          closestCentroidIndex = j;
        }
      }
 
      if (labels[i] !== closestCentroidIndex) {
        labels[i] = closestCentroidIndex;
        changed = true;
      }
    }
 
    // If no point changed its cluster assignment, we have converged
    if (!changed) {
      console.log(`Converged at iteration ${iter}`);
      break;
    }
 
    // 3. Update Step
    const sumsX = new Array(k).fill(0);
    const sumsY = new Array(k).fill(0);
    const counts = new Array(k).fill(0);
 
    for (let i = 0; i < points.length; i++) {
      const label = labels[i];
      sumsX[label] += points[i].x;
      sumsY[label] += points[i].y;
      counts[label]++;
    }
 
    for (let j = 0; j < k; j++) {
      if (counts[j] > 0) {
        centroids[j] = {
          x: sumsX[j] / counts[j],
          y: sumsY[j] / counts[j],
        };
      }
    }
  }
 
  return { centroids, labels };
}
 
// Test the example
const dataPoints: Point2D[] = [
  { x: 2, y: 1 },
  { x: 6, y: 5 },
  { x: 6, y: 1 },
  { x: 6, y: 3 },
  { x: 1, y: 1 },
];
 
const result = kmeans2D(dataPoints, 2);
console.log("Centroids:", result.centroids);
console.log("Labels:", result.labels);