np_kmeans.py in Python

Cluster 2-D points with Lloyd's k-means in NumPy.

"""Cluster 2-D points with a few rounds of Lloyd's algorithm."""
import numpy as np


def kmeans(points, k, iterations=10, seed=0):
    rng = np.random.default_rng(seed)
    centers = points[rng.choice(len(points), size=k, replace=False)]
    labels = np.zeros(len(points), dtype=int)
    for _ in range(iterations):
        distances = np.linalg.norm(
            points[:, np.newaxis] - centers[np.newaxis, :], axis=2)
        labels = distances.argmin(axis=1)
        for c in range(k):
            members = points[labels == c]
            if len(members):
                centers[c] = members.mean(axis=0)
    return centers, labels


def main():
    rng = np.random.default_rng(1)
    blob_a = rng.normal(loc=[0, 0], scale=0.5, size=(50, 2))
    blob_b = rng.normal(loc=[5, 5], scale=0.5, size=(50, 2))
    points = np.vstack([blob_a, blob_b])
    centers, labels = kmeans(points, k=2)
    for c, center in enumerate(centers):
        size = int((labels == c).sum())
        print(f"cluster {c}: {size} points near {center.round(2)}")


if __name__ == "__main__":
    main()

How it works

1. Broadcasting computes every point-to-center distance.
2. argmin assigns each point to its nearest center.
3. Centers move to the mean of their members each round.

Keywords and builtins used here

• as
• def
• enumerate
• for
• if
• int
• kmeans
• len
• main
• print
• range
• return

The run, in numbers

Lines

32

Characters to type

940

Tokens

303

Three-star pace

115 tpm

At the three-star pace of 115 tokens a minute, this run takes about 158 seconds.

Type this snippet

Step 1 of 1 in Encore, step 22 of 22 in NumPy.

← Previous