How to do clustering with K-Means in Python (scikit-learn)

Clustering with K-Means in Python lets you automatically group similar observations without needing labels. It is one of the most widely used unsupervised learning techniques in everyday work: it helps segment customers, organise products or explore hidden patterns in a dataset. With the scikit-learn library, applying it is quick and simple, as you will see below.

Prerequisites

• Python 3 installed on your computer.
• The scikit-learn and numpy libraries (install them with pip install scikit-learn numpy).
• Basic Python knowledge: lists, functions and loops.

Step 1: Prepare the data

Start by gathering the data you want to group. Each row is an observation and each column is a feature. In this example we use points with two features so they are easy to visualise, but K-Means works just as well with many more.

import numpy as np

# Each row is an observation, each column a feature
X = np.array([
    [1.0, 2.0],
    [1.5, 1.8],
    [5.0, 8.0],
    [8.0, 8.0],
    [1.0, 0.6],
    [9.0, 11.0],
])

Step 2: Scale the variables

K-Means relies on distances between points, so variables with very different scales can bias the groups. Use StandardScaler to bring every feature to the same scale before training the model.

from sklearn.preprocessing import StandardScaler

scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

Step 3: Choose the number of clusters

You need to decide how many groups (k) to split the data into. The elbow method helps with that choice: you train the model for several values of k and look at the inertia, that is, the sum of the distances from each point to the centre of its cluster. The point where the curve bends like an elbow usually indicates a good value of k.

from sklearn.cluster import KMeans

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

print(inertias)

When you run this code, you get a list of inertias, one per value of k. Look for the k from which the inertia stops dropping sharply: that is usually the most suitable number of clusters.

Step 4: Train the K-Means model

K-Means splits the data into k groups: it places each observation in the cluster whose centre is closest and keeps recomputing those centres until they stabilise. With your chosen value of k (here, 2), create and train the model. The n_init parameter sets how many times the algorithm runs with different initial centres, keeping the best result; random_state makes the result reproducible.

model = KMeans(n_clusters=2, random_state=42, n_init=10)
model.fit(X_scaled)

print(model.labels_)           # cluster of each observation
print(model.cluster_centers_)  # centre of each cluster

Printing model.labels_, each number tells you the cluster assigned to the observation in the same position; model.cluster_centers_, in turn, shows the coordinates of each group's centre, handy for understanding what each cluster represents.

Tip: K-Means assumes clusters that are roughly spherical and similar in size. If your groups have elongated shapes or very different densities, algorithms such as DBSCAN may give better results.

Step 5: Predict the cluster of new data

Once trained, the model can assign new points to one of the existing clusters. Remember to apply the same scaling to the new data before predicting, otherwise the result will not be reliable.

new_point = scaler.transform([[2.0, 2.0]])
print(model.predict(new_point))

Check the result

To confirm that the grouping makes sense, check that model.labels_ has a label (0, 1, ...) for each observation and that the points you expected to see together ended up in the same cluster. You can also compare the inertias from Step 3: the lower the inertia, the more compact the groups. Be aware, however, that always increasing the number of clusters lowers the inertia artificially, so follow the elbow method instead of picking the highest k.

Conclusion

You now have a complete workflow to do clustering with K-Means in Python: prepare the data, scale it, choose the number of clusters, train the model and predict new points. A good next step is to visualise the groups in a scatter plot with matplotlib, so you can see the clusters with your own eyes. So, which dataset will you try to group first?