Build K-Means With Python

Install scikit-learn

K-Means is available in the scikit-learn Python library.

pip install scikit-learn

If you are working inside a Jupyter Notebook, you can also use:

!pip install scikit-learn

Import KMeans

First, import the KMeans class.

from sklearn.cluster import KMeans

Now Python knows that we want to use the K-Means algorithm from scikit-learn.

Create Some Data

Let's create a small customer dataset.

We will use two features:

import numpy as np

X = np.array([
    [1000, 2],
    [1200, 3],
    [1100, 2],
    [10000, 15],
    [11000, 17],
    [10500, 16]
])

Each row represents one customer.

Choose the Number of Clusters

We learned earlier that K tells K-Means how many clusters to create.

Let's choose:

In Python:

model = KMeans(
    n_clusters=2,
    random_state=42
)

Here:

Train the K-Means Model

Now we give our data to the model.

model.fit(X)

This is where K-Means actually performs the clustering.

Get the Cluster Labels

After fitting the model, we can see which cluster each data point belongs to.

labels = model.labels_

print(labels)

You might get something similar to:

[1 1 1 0 0 0]

This means:

Get the Cluster Centers

K-Means also gives us the center of each cluster.

print(model.cluster_centers_)

These are called centroids.

Each centroid represents the average position of the data points in that cluster.

Build the Complete Model

Now let's put everything together.

import numpy as np
from sklearn.cluster import KMeans

# Customer data
X = np.array([
    [1000, 2],
    [1200, 3],
    [1100, 2],
    [10000, 15],
    [11000, 17],
    [10500, 16]
])

# Create the model
model = KMeans(
    n_clusters=2,
    random_state=42
)

# Train the model
model.fit(X)

# Get cluster labels
labels = model.labels_

# Get cluster centers
centers = model.cluster_centers_

print("Labels:")
print(labels)

print("Centers:")
print(centers)

What Did We Actually Do?

That's the basic K-Means workflow in Python.

Feature Scaling Can Matter

K-Means uses distance to decide which points are close to each other.

Therefore, features with very different scales can affect the result.

For example:

Spending has much larger numerical values than purchases. In a real machine-learning workflow, we would usually consider scaling the features before applying K-Means.

Check Your Understanding