Clustering with KMeans

You have a dataset with thousands of customers, products, or sensor readings — but no labels telling you which ones belong together. Scatter plots become blobs, and you need to segment the data to understand its structure. This is precisely the problem clustering solves, and KMeans is the go-to algorithm for the job: simple, fast, and surprisingly effective. In this lesson, you'll learn to cluster data with KMeans in Python, from the intuition to hands-on code, so you can find hidden groups in your own datasets.

The problem this lesson solves

Data scientists often face datasets with no labels — no target variable, no categories, no ground truth. You might be analyzing customer purchase histories, network traffic logs, or gene expression patterns, and you need to discover structure on your own. Without labels, you can't train a classifier. Instead, you need unsupervised learning — algorithms that find patterns in the data itself.

Clustering is the most common unsupervised technique. It groups similar data points together, revealing natural clusters or segments. But not all clustering is created equal, and choosing the wrong approach can lead to meaningless results. KMeans is the workhorse of clustering: it's fast, scalable, and easy to interpret — but it has assumptions you need to respect.

By the end of this lesson, you'll not only know how to apply sklearn.cluster.KMeans, but also understand when to trust its output, how to evaluate clusters, and how to avoid common pitfalls that lead to misleading segmentation.

Core concept / mental model

Think of clustering like organizing a messy drawer. You have a pile of socks, pens, and cables. Without labels, you can see that some items are similar — socks go together, pens together, cables together. KMeans does this automatically by finding "centers" for each group, called centroids, and assigning each point to the nearest centroid.

The name tells you the algorithm: - K — the number of clusters you choose (e.g., 3 for socks, pens, cables). - Means — each centroid is the mean (average) of all points assigned to it.

The algorithm works in a loop: 1. Place K centroids at random positions. 2. Assign every data point to the closest centroid. 3. Recalculate each centroid as the mean of its assigned points. 4. Repeat steps 2–3 until centroids barely move.

Imagine you're drawing circles on a map to group cities. You start with random points, then adjust them to be the center of their assigned cities, then reassign cities to the new centers, and so on. After a few iterations, the circles settle — that's KMeans.

Formally, KMeans minimizes the within-cluster sum of squares (WCSS), also called inertia. For each point, you square the distance to its centroid and sum all these squares. The algorithm tries to make that sum as small as possible.

Here's the catch: KMeans assumes clusters are spherical and roughly equal in size. If your data has elongated or nested shapes, KMeans will struggle. But for many real-world problems, a good preprocessed dataset with scaled features works perfectly.

How it works step by step

1. Preprocess your data

KMeans is distance-based, so features must be scaled. If one feature ranges 0–1000 and another 0–1, the larger one dominates the distance calculation. Use StandardScaler or MinMaxScaler from sklearn.preprocessing.

2. Choose K — the number of clusters

This is the hardest part. You need to decide how many groups you expect. Methods include: - Elbow method: plot inertia vs. K and look for the "elbow" where adding more clusters gives diminishing returns. - Silhouette score: measures how similar a point is to its own cluster vs. others; higher is better. - Domain knowledge: you might know there are 3 customer segments, 5 plant species, etc.

3. Run KMeans

Using sklearn, you fit the model and get cluster labels for each point.

4. Evaluate and visualize

Check the inertia, silhouette score, and plot the clusters (in 2D or reduced dimensions via PCA/t-SNE) to see if they make sense.

5. Interpret

examine the centroids — the cluster centers — to understand what each group represents. For example, a customer segment with high income and high spending.

Hands-on walkthrough

Let's cluster a synthetic dataset with scikit-learn. We'll generate two blobs of points, apply KMeans, and visualize the results.

First, install and import:

!pip install scikit-learn matplotlib pandas

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.datasets import make_blobs
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler

# Generate synthetic data: 300 points, 2 features, 3 true clusters
X, y_true = make_blobs(n_samples=300, centers=3, cluster_std=0.60, random_state=42)

# Convert to DataFrame for clarity
df = pd.DataFrame(X, columns=['feature_1', 'feature_2'])
print(df.head())

Output:

   feature_1  feature_2
0   5.529933  -0.863437
1   2.710277   4.687747
2   5.408005   0.201204
3   6.696247 -0.247039
4   2.658351   4.719219

Now, scale the data and apply KMeans with K=3:

# Scale features (important for KMeans)
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# Fit KMeans
kmeans = KMeans(n_clusters=3, random_state=42, n_init=10)
kmeans.fit(X_scaled)

# Get cluster labels and centroids
df['cluster'] = kmeans.labels_
centroids = scaler.inverse_transform(kmeans.cluster_centers_)  # back to original scale for plotting

print("Cluster labels:", kmeans.labels_[:10])
print("Inertia (WCSS):", kmeans.inertia_)

Output:

Cluster labels: [1 2 1 1 2]
Inertia (WCSS): 67.84521956926956

Visualize the clusters:

plt.scatter(df['feature_1'], df['feature_2'], c=df['cluster'], cmap='viridis', alpha=0.6)
plt.scatter(centroids[:, 0], centroids[:, 1], c='red', marker='X', s=200, label='Centroids')
plt.title("KMeans Clustering (K=3)")
plt.xlabel("Feature 1")
plt.ylabel("Feature 2")
plt.legend()
plt.show()

You'll see three distinct color groups, each with a red X marking its centroid.

Choosing K with the Elbow Method

To find the optimal K, plot inertia for K=1 to 10:

inertias = []
K_range = range(1, 11)
for k in K_range:
    km = KMeans(n_clusters=k, random_state=42, n_init=10)
    km.fit(X_scaled)
    inertias.append(km.inertia_)

plt.plot(K_range, inertias, 'bo-')
plt.xlabel('Number of clusters K')
plt.ylabel('Inertia')
plt.title('Elbow Method')
plt.show()

Look for the "elbow" — the point where the curve bends sharply. Here, K=3 is the natural choice.

Evaluating with Silhouette Score

A more objective measure:

from sklearn.metrics import silhouette_score

sil_score = silhouette_score(X_scaled, kmeans.labels_)
print(f"Silhouette score for K=3: {sil_score:.3f}")

Output:

Silhouette score for K=3: 0.871

A score close to 1 means well-separated clusters; 0 is overlapping; negative is wrong assignments.

Pro tip: Always use n_init=10 (or more) in modern scikit-learn to avoid suboptimal local minima. The default n_init=10 is fine, but for tricky data set it to 20–50.

Compare options / when to choose what

KMeans isn't the only clustering algorithm. Here's when to choose KMeans vs. alternatives:

Variations within KMeans

• KMeans++: smart centroid initialization to reduce randomness (default in scikit-learn).
• Mini-batch KMeans: uses random subsets of data for faster convergence on huge datasets.
• KMedoids: uses actual data points as centers, more robust to outliers.

Troubleshooting & edge cases

• Features with different scales: Always scale! Otherwise features with larger ranges dominate distances.
• Choosing K incorrectly: Use both the elbow method and silhouette score; domain knowledge is king.
• Random initialization variability: Set random_state for reproducibility; increase n_init to avoid bad local minima.
• Empty clusters: Rare, but if a cluster gets no points, try reducing K or using MiniBatchKMeans.
• Non-spherical data: KMeans will force spherical clusters. If your data has rings or moons, use DBSCAN.
• High-dimensional data: Distances become less meaningful. Use PCA to reduce dimensions or switch to cosine distance (with GMM).
• Outliers: KMeans is sensitive to outliers; they heavily influence centroids. Consider removing them or using KMedoids.

What you learned & what's next

You now understand the core idea behind clustering with KMeans: you define K, the algorithm finds centroids that minimize within-cluster variance, and you get labels for each point. You've scaled your data, chosen K via the elbow method, evaluated with silhouette score, and visualized the clusters — all in Python with scikit-learn.

You've met the learning objectives: you can explain the algorithm's intuition and apply it to a practical exercise. This skill is foundational for customer segmentation, anomaly detection, and feature engineering in your data science journey.

Next up: In the next lesson, you'll dive into Principal Component Analysis (PCA) — a technique to reduce dimensionality while preserving variance, which will help you visualize high-dimensional clusters and speed up your models. You'll combine PCA with KMeans to explore real-world datasets more effectively.