AI360Xpert
Coding & Implementation

Coding & Implementation

2 interview questions in this topic. Expand any question to read the full answer.

Implement k-means clustering from scratch

Answer

K-means partitions points into k clusters by minimizing the total squared distance from each point to its cluster's center (the inertia). It uses an alternating optimization (Lloyd's algorithm): pick k initial centroids, then loop two steps until nothing moves — assign each point to its nearest centroid, then update each centroid to the mean of the points assigned to it. Each iteration provably lowers the inertia, so it converges, but only to a local optimum.

import numpy as np
def kmeans(X: np.ndarray, k: int, iters: int = 100) -> tuple[np.ndarray, np.ndarray]:    """Cluster X (n×d) into k groups; return centroids and integer labels."""    rng = np.random.default_rng(0)    centroids = X[rng.choice(len(X), k, replace=False)]  # random distinct seeds    for _ in range(iters):        # Assign: squared distance from every point to every centroid → nearest.        dists = ((X[:, None, :] - centroids[None, :, :]) ** 2).sum(axis=2)        labels = dists.argmin(axis=1)        # Update: each centroid becomes the mean of its assigned points.        new = np.array([X[labels == j].mean(axis=0) for j in range(k)])        if np.allclose(new, centroids):  # converged — assignments stable            break        centroids = new    return centroids, labels

The assignment step broadcasts to an (n, k) distance matrix and takes the argmin per point; the update step averages each cluster; allclose is the convergence test. The things to flag are the limitations: it finds only a local optimum, so initialization matters — use k-means++ to spread the initial centroids out (and often run it several times, keeping the lowest inertia). k isn't learned, so pick it with the elbow method or silhouette score. And it assumes roughly spherical, similar-sized clusters and uses Euclidean distance, so standardize features first and switch to DBSCAN or a Gaussian mixture when clusters aren't spherical.

K-means alternates two steps until convergence: assign each point to the nearest centroid, then move each centroid to the mean of its assigned points; clusters and centroids stabilize over iterations.
K-means alternates two steps until convergence: assign each point to the nearest centroid, then move each centroid to the mean of its assigned points; clusters and centroids stabilize over iterations.

💡 Note Imagine k food trucks deciding where to park in a city. Each resident walks to the nearest truck (assignment), then each truck repositions to the middle of its own customers (update). Residents re-pick, trucks move again, and it repeats until nobody switches and no truck needs to move. That settling point is the clustering.

Related Questions

Implement softmax from scratch

Answer

Softmax converts a vector of raw scores (logits) into a probability distribution — every output is between 0 and 1 and they sum to 1. For an input vector x:

softmax(x)ᵢ = exp(xᵢ) / Σⱼ exp(xⱼ)

Exponentiating guarantees positive outputs; dividing by the sum makes them sum to 1. It's also monotonic, and the exponential exaggerates gaps, so a clear winner gets a confidently high probability.

The thing that makes or breaks the implementation is numerical stability. Taken literally, exp(1000) overflows to infinity and you get inf/inf = NaN. The fix uses a simple identity: subtracting a constant c from every input leaves softmax unchanged, because the exp(-c) factor cancels top and bottom. Choose c = max(x) so the largest shifted value is 0 and the biggest exponent is exp(0) = 1 — no overflow.

import numpy as np
def softmax(x: np.ndarray) -> np.ndarray:    """Numerically stable softmax over the last axis of x."""    # Subtract the max (per row) so the largest exponent is exp(0) = 1 — no overflow.    shifted = x - np.max(x, axis=-1, keepdims=True)    exp = np.exp(shifted)                       # all values in (0, 1]    return exp / np.sum(exp, axis=-1, keepdims=True)  # normalize to sum to 1

The axis=-1 and keepdims=True make it work row-wise on a batch, which is the usual (batch, classes) case. Two related points worth raising: for log-probabilities you'd use log-sum-exp / log_softmax (what cross-entropy uses under the hood), and dividing the logits by a temperature before softmax is the knob behind LLM sampling — high T flattens the distribution, low T sharpens it.

Softmax maps a vector of raw logits to a probability distribution: exponentiate each score, divide by the sum of exponentials, and the largest logit becomes the largest probability while all outputs sum to one.
Softmax maps a vector of raw logits to a probability distribution: exponentiate each score, divide by the sum of exponentials, and the largest logit becomes the largest probability while all outputs sum to one.

💡 Note Softmax is like tallying votes into percentages, but with a megaphone for the front-runner: a candidate slightly ahead in raw score ends up comfortably ahead in probability. Subtracting the max is just measuring everyone's score relative to the leader first — it changes the numbers you plug in, not the final percentages.

Related Questions