Menu
CoddyTech

K-Means in scikit-learn

Part of the Unsupervised Learning section of Coddy's Machine Learning journey. Lesson 4 of 51.

scikit-learn runs the whole k-means loop for you, here on 180 apples scored for sweetness and crunch:

180 apples, sweetness across and crunch up, in three groups with a black star at each center: sweet and crunchy at the upper right, soft at the bottom, tart and crunchy at the upper left. A new apple at the top right joins cluster 0

Fit with X alone, no y:

from sklearn.cluster import KMeans
km = KMeans(n_clusters=3, n_init=10, random_state=0)
km.fit(X)
print(km.labels_[:8])
print(km.cluster_centers_.round(1))

This prints 8 rows' cluster numbers, then one row per center:

[2 0 0 1 2 0 1 1]
[[7.  7.5]
 [6.  3. ]
 [2.9 8.1]]

n_init=10 tries 10 starts and keeps the tightest. km.predict([[8.6, 9.2]]) gives the new apple its nearest center, [0].

labels = km.fit_predict(X) does both at once: it fits, then returns labels_.

The numbers are only names: on other data, cluster 0 can be any group.

challenge icon

Challenge

Easy

A biologist logged 150 dives of one humpback whale: dive time and the surface rest after it. Are there kinds of dive? The starter sets up the chart.

  1. Fit KMeans on X: 3 clusters, n_init=10, random_state=0.
  2. Scatter each cluster as its own group, labelled cluster 0 to cluster 2.
  3. Mark the centers with black stars, labelled centers. Add a legend and show.

Like this, from another whale:

A scatter titled Whale Dives, dive minutes across and surface rest up: three clusters, short dives with long rests at the upper left, short dives with short rests at the bottom, long dives at the right, a black star in each, and a legend

Coddy describes your chart.

Try it yourself

import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans

# Don't change the lines below
rng = np.random.default_rng(14)
dive = np.concatenate([rng.normal(4, 1.0, 60), rng.normal(12, 2.0, 50), rng.normal(2.5, 0.7, 40)])
rest = np.concatenate([rng.normal(1.5, 0.5, 60), rng.normal(4, 1.0, 50), rng.normal(8, 1.5, 40)])
# 150 dives of one humpback whale: dive time and the surface rest after it, in minutes
X = np.round(np.clip(np.column_stack([dive, rest]), 0.5, None), 1)

fig, ax = plt.subplots()
ax.set_title("Whale Dives")
ax.set_xlabel("Dive (min)")
ax.set_ylabel("Surface rest (min)")

# Write your code below
quiz iconTest yourself

This lesson includes a short quiz. Start the lesson to answer it and track your progress.

All lessons in Unsupervised Learning

Practice on your own: Online Python compiler