K-Means Clustering

K-Means is an unsupervised learning algorithm for clustering that divides data into k groups based on the distance to a centroid. It is popular because it is simple and efficient.

How It Works

1. Choose k centroids randomly
2. Assign each data point to the nearest centroid
3. Recompute each centroid as the mean of all points in the cluster
4. Repeat steps 2-3 until the centroids stop changing (convergence)

from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs
import numpy as np

X, y_true = make_blobs(n_samples=300, centers=4, random_state=42)

kmeans = KMeans(n_clusters=4, random_state=42, n_init=10)
kmeans.fit(X)

labels = kmeans.labels_
centroids = kmeans.cluster_centers_

print(f"Cluster labels: {np.unique(labels)}")
print(f"Centroid shape: {centroids.shape}")

Elbow Method

The elbow method helps determine the optimal number of clusters by examining the total within-cluster sum of squares (WCSS) for various values of k.

from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs
import matplotlib.pyplot as plt

X, _ = make_blobs(n_samples=300, centers=4, random_state=42)

wcss = []
k_range = range(1, 11)

for k in k_range:
    kmeans = KMeans(n_clusters=k, random_state=42, n_init=10)
    kmeans.fit(X)
    wcss.append(kmeans.inertia_)

plt.figure(figsize=(8, 4))
plt.plot(k_range, wcss, 'bo-')
plt.xlabel('Number of Clusters (k)')
plt.ylabel('WCSS (Inertia)')
plt.title('Elbow Method')
plt.grid(True, alpha=0.3)
plt.show()

Silhouette Score

The silhouette score measures how similar a data point is to its own cluster compared to other clusters. The value ranges from -1 to 1; higher is better.

from sklearn.metrics import silhouette_score
from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs

X, _ = make_blobs(n_samples=300, centers=4, random_state=42)

for k in range(2, 8):
    kmeans = KMeans(n_clusters=k, random_state=42, n_init=10)
    labels = kmeans.fit_predict(X)
    score = silhouette_score(X, labels)
    print(f"k={k}: Silhouette Score = {score:.4f}")

K-Means++ Initialization

K-Means++ chooses initial centroids that are well spread out, reducing convergence to a local optimum. This is the default in scikit-learn.

from sklearn.cluster import KMeans

# The default already uses k-means++
kmeans = KMeans(
    n_clusters=3,
    init='k-means++',
    n_init=10,
    random_state=42
)

Important Attributes

from sklearn.cluster import KMeans
from sklearn.datasets import load_iris

iris = load_iris()
kmeans = KMeans(n_clusters=3, random_state=42, n_init=10)
kmeans.fit(iris.data)

print(f"Labels: {kmeans.labels_}")
print(f"Centroids:\n{kmeans.cluster_centers_}")
print(f"Inertia (WCSS): {kmeans.inertia_:.4f}")
print(f"Number of iterations: {kmeans.n_iter_}")

Example: Customer Segmentation

import pandas as pd
import numpy as np
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler

np.random.seed(42)
data = {
    'annual_income': np.random.randint(20000, 150000, 200),
    'spending_score': np.random.randint(1, 100, 200)
}
df = pd.DataFrame(data)

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

kmeans = KMeans(n_clusters=4, random_state=42, n_init=10)
df['cluster'] = kmeans.fit_predict(X_scaled)

cluster_summary = df.groupby('cluster').agg({
    'annual_income': 'mean',
    'spending_score': 'mean'
}).round(0)
print(cluster_summary)

Pros and Cons

Next chapter: Hierarchical Clustering, an alternative approach to grouping data.

Next: Hierarchical Clustering →