K-Means++ Initialization

Random initialization is the weakest part of Lloyd’s k-means: seed two centroids inside the same blob and the algorithm can converge to a bad local optimum. K-means++ fixes this by spreading the initial centroids out. It picks the first center at random, then biases every later pick toward points far from the centers already chosen. The interview task is to turn that into vectorized NumPy plus a correct weighted sample.

Definition

Let D(x) be the Euclidean distance from a point x to the nearest centroid chosen so far. The seeding procedure is:

1. Choose the first centroid c_1 uniformly at random from the rows of X.
2. For each remaining centroid:
     - compute D(x)^2 for every point x (distance to its nearest chosen centroid)
     - choose the next centroid to be point x with probability D(x)^2 / sum_x D(x)^2
3. Stop once k centroids have been chosen.

The weighting by D(x)^2 (squared, not raw distance) is the defining detail: points that are already close to a chosen center get near-zero weight, while far-away points dominate the draw.

Task

Complete kmeans_plus_plus_init(X, k, seed) so it returns the (k, d) array of chosen centroids in pick order. Create one generator with np.random.default_rng(seed) and use it for both the first uniform pick and every subsequent weighted pick. Maintain a running vector of each point’s squared distance to its nearest chosen centroid, updating it with an element-wise minimum after each new centroid rather than recomputing against all centroids. Do not call sklearn, scipy, or any built-in k-means or sampling helper beyond the NumPy generator.

Example

X = np.array([[0.0,  0.0],
              [0.0,  1.0],
              [1.0,  0.0],
              [10.0, 10.0],
              [10.0, 11.0],
              [11.0, 10.0]])
kmeans_plus_plus_init(X, k=2, seed=0)
# -> array([[11., 10.],
#           [ 0.,  0.]])

With seed=0 the first uniform draw lands on row 5, [11, 10], in the top-right blob. Every point in that blob then has a tiny squared distance and near-zero weight, so the second pick is drawn almost entirely from the far blob and lands on [0, 0] — the two seeds sit in different clusters, which is the whole point of k-means++.

Constraints

• 1 <= k <= n, X has shape (n, d) with 1 <= n <= 10^4, 1 <= d <= 10^3.
• Use the passed-in seed through a single np.random.default_rng(seed); the same seed must always return the same centroids.
• Each returned centroid must be an actual row of X, and the result is returned in the order the centroids were picked.
• Weight the draw by squared distance, and keep the running nearest-distance vector via an element-wise minimum (do not recompute distances to already-chosen centroids).

Approach

Draw the first centroid uniformly with a seeded generator, then maintain a length-n vector closest_sq holding each point’s squared distance to its nearest chosen centroid. Each subsequent centroid is sampled with probability proportional to closest_sq, so far-away points dominate the draw. After every pick, refresh closest_sq with an element-wise minimum against the squared distances to the new centroid — that keeps each round O(n*d) instead of recomputing against every chosen centroid.

Solution

import numpy as np

def kmeans_plus_plus_init(X: np.ndarray, k: int, seed: int = 0) -> np.ndarray:
    X = np.asarray(X, dtype=np.float64)
    n, d = X.shape
    rng = np.random.default_rng(seed)

    centroids = np.empty((k, d), dtype=np.float64)

    # 1. first centroid: uniform over the rows
    first = int(rng.integers(n))
    centroids[0] = X[first]

    # squared distance from every point to its nearest chosen centroid
    closest_sq = np.sum((X - centroids[0]) ** 2, axis=1)

    for i in range(1, k):
        total = closest_sq.sum()
        if total == 0.0:
            # all points coincide with a chosen centroid; fall back to uniform
            probs = np.full(n, 1.0 / n)
        else:
            probs = closest_sq / total

        # 2. weighted draw: P(x) proportional to D(x)^2
        idx = int(rng.choice(n, p=probs))
        centroids[i] = X[idx]

        # 3. update the running nearest-distance with an element-wise min
        new_sq = np.sum((X - centroids[i]) ** 2, axis=1)
        closest_sq = np.minimum(closest_sq, new_sq)

    return centroids

Walkthrough

On the example, X has two blobs and seed=0, k=2.

• rng.integers(6) returns 5, so centroids[0] = X[5] = [11, 10].
• closest_sq = squared distance of each row to [11, 10]: [221, 202, 200, 1, 2, 0]. The three top-right points (rows 3-5) are tiny; the three bottom-left points (rows 0-2) are large.
• total = 626, so probs ≈ [0.353, 0.323, 0.319, 0.0016, 0.0032, 0.0]. Almost all the mass sits on the far blob, so rng.choice draws from rows 0-2 and returns index 0.
• centroids[1] = X[0] = [0, 0]. The two seeds land in different clusters.
• The loop ends with k = 2 centroids, returned as [[11, 10], [0, 0]].

If instead the second draw had picked, say, row 1, the np.minimum update would then set each point’s closest_sq to the smaller of its distance to [11,10] and to the new center — the mechanism that lets a third pick avoid both existing seeds.

Complexity & notes

• Time O(knd): each of the k rounds computes one (n, d) squared-distance pass and one O(n) sample. Space O(n*d) for X plus O(n) for closest_sq and probs.
• The running np.minimum update is the key optimization. Recomputing each point’s nearest centroid over all i chosen centroids every round would be O(k^2nd); carrying closest_sq forward drops a full factor of k.
• Sampling weight must be the squared distance. Weighting by raw D(x) is a different (worse) scheme; the O(log k) competitive guarantee of k-means++ relies on D(x)^2.
• Use one generator seeded once. Reseeding inside the loop, or mixing np.random.* global calls with the generator, breaks reproducibility across the picks.
• The total == 0 guard matters when there are duplicate points or when k exceeds the number of distinct rows: once every point coincides with some chosen centroid, all weights are zero and rng.choice would raise on a zero-sum probability vector, so we fall back to a uniform draw.
• This returns only the seeds; a full k-means would follow with Lloyd iterations (assign, then recompute means) until the assignments stop changing.