K-Means from Scratch (k-means++ Init, Convergence, Empty Clusters) - Unsupervised Learning

K-Means from Scratch (k-means++ Init, Convergence, Empty Clusters)

What you'll build

The complete k-means algorithm (Lloyd's algorithm): k-means++ seeding, alternating assignment and update steps, a convergence test, and a rule for clusters that end up empty. Two functions:

• kmeans_pp_init(X, k, seed=0): k-means++ seeding, returns the (k, d) initial centroids
• kmeans(X, k, init='k-means++', max_iter=100, tol=1e-4, seed=0): the full algorithm

Background

k-means minimises the inertia sum_i ||x_i - C[label_i]||^2. Each assignment step and each update step can only lower it, so the algorithm converges, but only to a local minimum: initialisation matters. k-means++ spreads the initial centroids out by sampling each new one with probability proportional to its squared distance from the centroids chosen so far, which gives an O(log k) approximation guarantee in expectation.

k-means++ seeding (follow exactly, so results are reproducible)

Create one generator rng = np.random.default_rng(seed) and draw every random number with u = rng.random():

1. First centroid: the point at index floor(u * n).
2. Each further centroid: compute D2[i], the squared distance from point i to its nearest already chosen centroid, and cdf = np.cumsum(D2). Draw u and pick the first index i with cdf[i] > u * cdf[-1].
3. If cdf[-1] == 0 (every point coincides with a chosen centroid), pick index floor(u * n) instead (using the u already drawn for this centroid).

Lloyd iterations

Start from the initial centroids C (k-means++ when init == 'k-means++', otherwise init is a (k, d) array). For it = 1, 2, ..., max_iter:

1. Assign: each point goes to its nearest centroid by squared Euclidean distance; ties go to the lowest centroid index.
2. Update: each non-empty cluster's centroid becomes the mean of its points.
3. Empty clusters: a cluster with no points is moved onto the point farthest from its assigned centroid in this iteration's assignment (largest squared distance; ties go to the lowest point index). Handle several empty clusters in increasing cluster index; each point can be used only once.
4. Stop if the labels did not change since the previous iteration, or if the largest centroid shift max_j ||C_new[j] - C[j]|| is below tol.

Inputs

• X: data, shape (n, d); k: number of clusters, 1 <= k <= n
• init: 'k-means++' or a (k, d) array of initial centroids
• max_iter (int >= 1), tol (absolute shift threshold), seed (int, used only by k-means++)

Expected Return Format

kmeans_pp_init returns a (k, d) float array. kmeans returns a dictionary with these keys:

• 'centroids': (k, d) float array, the centroids after the last update
• 'labels': (n,) int array, each point's nearest returned centroid (one final assignment step against the returned centroids)
• 'inertia': Python float, sum_i ||x_i - centroids[labels_i]||^2 for those labels
• 'n_iter': int, the number of assign/update iterations performed (1 <= n_iter <= max_iter)

Differences from scikit-learn's KMeans

scikit-learn's k-means++ tries several candidates per step and keeps the best ("greedy k-means++"), scales tol by the data variance, and runs several initialisations (n_init). This problem uses the plain version so every number is reproducible.

Constraints

• NumPy only; vectorise the distance computations.

Follow-up

• How do you choose k? (Elbow on inertia, silhouette score.)
• What is mini-batch k-means, and when would you use it?

Input: kmeans(np.array([[1.0, 1.0], [1.5, 2.0], [3.0, 4.0], [5.0, 7.0], [3.5, 5.0], [4.5, 5.0], [3.5, 4.5]]), 2, init=np.array([[1.0, 1.0], [5.0, 7.0]]))
Output: {'centroids': [[1.25, 1.5], [3.9, 5.1]], 'labels': [0, 0, 1, 1, 1, 1, 1], 'inertia': 8.525, 'n_iter': 3}

Input: kmeans_pp_init(np.array([[0.0, 0.0], [0.0, 1.0], [10.0, 10.0], [10.0, 11.0], [20.0, 0.0]]), 3, seed=0)
Output: array([[10., 11.], [0., 0.], [20., 0.]])

import numpy as np

def kmeans_pp_init(X, k, seed=0):
    """
    k-means++ seeding (see the description for the exact sampling rule).

    Args:
        X: Data, array of shape (n, d)
        k: Number of clusters
        seed: Seed for np.random.default_rng

    Returns:
        centroids: Initial centroids, array of shape (k, d)
    """
    # Your code here
    pass

def kmeans(X, k, init='k-means++', max_iter=100, tol=1e-4, seed=0):
    """
    Full k-means (Lloyd's algorithm).

    Args:
        X: Data, array of shape (n, d)
        k: Number of clusters
        init: 'k-means++' or an initial centroid array of shape (k, d)
        max_iter: Maximum number of assign/update iterations
        tol: Stop when the largest centroid shift is below tol
        seed: Seed for k-means++

    Returns:
        dict with 'centroids' (k, d), 'labels' (n,), 'inertia' (float), 'n_iter' (int)
    """
    # Your code here
    pass