Implement K-Means Clustering - Lloyd's Algorithm from Scratch

CODINGintermediateInterview freq: very_high

Choose your preparation mode6 modes available

Learn Trace Debug Trade-off Challenge

🤖

Implement K-Means Clustering using Lloyd's Algorithm without sklearn.

codingintermediateFreq: very_highClassical ML Coding

Prerequisites:K-Means Clustering - InitializationConvergence

📝

Problem Statement

Implement K-Means Clustering using Lloyd's Algorithm from scratch. Given a dataset X of shape (n_samples, n_features) and an integer K, partition the dataset into K clusters by minimizing the within-cluster sum of squares.

I/O Specification

Input:  numpy.ndarray X, int K
Output: numpy.ndarray centroids, numpy.ndarray labels

Example

Input: X = np.array([[1, 2], [1, 4], [1, 0], [10, 2], [10, 4], [10, 0]]), K = 2 → Output: centroids = np.array([[1, 2], [10, 2]]), labels = np.array([0, 0, 0, 1, 1, 1])

🎯

Interview Angles

Testing understanding of clustering and initialization strategies.
Evaluating ability to optimize naive algorithms using vectorization.
Assessing handling of edge cases and numerical stability in iterative algorithms.
Understanding of convergence criteria and stopping conditions in iterative methods.

🔵

Naive Approach

Randomly initialize K centroids, then iteratively assign each point to the nearest centroid and update centroids as the mean of assigned points. Repeat until convergence.

Time: O(n * K * max_iters * d) - n samples, K centroids, d dimensions, max_iters iterations. Space: O(n + K * d) - storing labels and centroids.

Why Insufficient

O(n * K * max_iters * d) because it recomputes distances for each point to each centroid in every iteration - inefficient for large n or K.

🐍

Naive Code

Python

import numpy as np

def k_means_naive(X, K, max_iters=100):
    n_samples, n_features = X.shape
    centroids = X[np.random.choice(n_samples, K, replace=False)]
    for _ in range(max_iters):
        labels = np.argmin(np.linalg.norm(X[:, np.newaxis] - centroids, axis=2), axis=1)
        new_centroids = np.array([X[labels == k].mean(axis=0) for k in range(K)])
        if np.all(centroids == new_centroids):
            break
        centroids = new_centroids
    return centroids, labels

# Driver example
X = np.array([[1, 2], [1, 4], [1, 0], [10, 2], [10, 4], [10, 0]])
K = 2
centroids, labels = k_means_naive(X, K)
print('Centroids:', centroids)
print('Labels:', labels)

💡

Key Insight

Use vectorized operations to compute distances and update centroids, reducing the need for explicit loops and improving computational efficiency.

🐍

Optimized Code

Python

import numpy as np

def k_means_optimized(X, K, max_iters=100):
    n_samples, n_features = X.shape
    centroids = X[np.random.choice(n_samples, K, replace=False)]
    for _ in range(max_iters):
        distances = np.linalg.norm(X[:, np.newaxis] - centroids, axis=2)
        labels = np.argmin(distances, axis=1)
        new_centroids = np.array([X[labels == k].mean(axis=0) for k in range(K)])
        if np.allclose(centroids, new_centroids):
            break
        centroids = new_centroids
    return centroids, labels

# Driver example
X = np.array([[1, 2], [1, 4], [1, 0], [10, 2], [10, 4], [10, 0]])
K = 2
centroids, labels = k_means_optimized(X, K)
print('Centroids:', centroids)
print('Labels:', labels)

🔍

Line-by-Line

centroids = X[np.random.choice(n_samples, K, replace=False)]

Randomly select K initial centroids from the dataset to start the clustering process.

distances = np.linalg.norm(X[:, np.newaxis] - centroids, axis=2)

Compute the Euclidean distance from each point to each centroid using vectorized operations.

labels = np.argmin(distances, axis=1)

Assign each point to the nearest centroid by finding the index of the minimum distance.

new_centroids = np.array([X[labels == k].mean(axis=0) for k in range(K)])

Update centroids by computing the mean of all points assigned to each centroid.

if np.allclose(centroids, new_centroids)

Check for convergence by comparing old and new centroids with a tolerance for numerical stability.

centroids = new_centroids

Update centroids for the next iteration if convergence is not reached.

⚡

Complexity

Time: O(n * K * d * max_iters) - n samples, K centroids, d dimensions, max_iters iterations. Space: O(n + K * d) - storing labels and centroids.

Output

Centroids: [[ 1.  2.]
 [10.  2.]]
Labels: [0 0 0 1 1 1]

🔲

Edge Cases

Empty Dataset

The function should handle an empty dataset gracefully without errors.

n_samples, n_features = X.shape

Single Element Dataset

The function should return the single point as the centroid if K=1.

centroids = X[np.random.choice(n_samples, K, replace=False)]

All Points Identical

The function should converge immediately as all points are already in one cluster.

if np.allclose(centroids, new_centroids)

⚠️

Common Mistakes

Incorrect centroid initialization leading to poor convergence.

Symptom: Algorithm takes too long to converge or converges to suboptimal clusters.

Not handling empty clusters during centroid update.

Symptom: Runtime errors or incorrect centroid updates when a cluster has no points.

Using a fixed number of iterations without checking for convergence.

Symptom: Unnecessary computations leading to inefficiency.

Incorrect distance computation leading to wrong cluster assignments.

Symptom: Points are assigned to incorrect clusters, resulting in poor clustering.

Practice

(1/5)

1. What is the time complexity of calculating the mean Average Precision (mAP) for a single image with N predicted bounding boxes and M ground truth boxes?

easy

A. O(N + M)

B. O(N^2)

C. O(N * M)

D. O(M^2)

Solution

Step 1: Understand the Process
For each predicted box, IoU is calculated with each ground truth box, leading to N * M comparisons.
Step 2: Determine Complexity
Each comparison is constant time, so the overall complexity is O(N * M).
Final Answer:
Option C -> Option C
Quick Check:
IoU comparisons drive complexity [OK]

Hint: IoU comparisons drive complexity [OK]

Common Mistakes:

2. What is the time complexity of generating a sinusoidal positional encoding matrix of size position x d_model?

easy

A. O(d_model^2)

B. O(position^2)

C. O(position * d_model)

D. O(1)

Solution

Step 1: Analyze the operations
The code iterates over each position and each dimension, performing constant-time operations.
Step 2: Eliminate wrong options
O(position^2) is incorrect because the complexity is not quadratic in position. O(d_model^2) is incorrect because the complexity is not quadratic in d_model. O(1) is incorrect because the complexity is not constant.
Final Answer:
Option C -> Option C
Quick Check:
Iterate over position and dimension. [OK]

Hint: Iterate over position and dimension. [OK]

Common Mistakes:

3. What is the time complexity of computing the scaled dot-product attention for a single query with n keys?

easy

A. O(n^2)

B. O(n log n)

C. O(n)

D. O(1)

Solution

Step 1: Identify operations
For a single query, the main operations are dot product and softmax, both O(n).
Step 2: Eliminate incorrect complexities
O(n^2) and O(n log n) are incorrect as they imply more complex operations than present.
Final Answer:
Option C -> Option C
Quick Check:
Single query, n keys -> O(n) [OK]

Hint: Single query, n keys -> O(n) [OK]

Common Mistakes:

4. In a recommendation system, why might you choose cosine similarity over Euclidean distance?

medium

A. Cosine similarity is invariant to vector magnitude.

B. Cosine similarity is more interpretable than Euclidean distance.

C. Cosine similarity handles missing data better.

D. Cosine similarity is faster to compute for large vectors.

Solution

Step 1: Compare properties
Cosine similarity measures angle, not magnitude, making it invariant to vector length.
Step 2: Eliminate incorrect reasons
Cosine similarity isn't inherently faster (A), doesn't handle missing data better (C), and interpretability is subjective (D).
Final Answer:
Option A -> Option A
Quick Check:
Cosine ignores magnitude, focuses on direction. [OK]

Hint: Cosine ignores magnitude, focuses on direction. [OK]

Common Mistakes:

5. How can you optimize a KD-Tree implementation for a KNN classifier to handle large datasets more efficiently?

hard

A. Increase the number of dimensions.

B. Use approximate nearest neighbors.

C. Use a balanced KD-Tree.

D. Reduce the number of data points.

Solution

Step 1: Identify optimization strategies
Approximate nearest neighbors can significantly speed up queries in large datasets.
Step 2: Evaluate other options
A: Balancing helps but isn't enough for large datasets. B: Increasing dimensions worsens performance. D: Reducing data points isn't always feasible.
Final Answer:
Option B -> Option B
Quick Check:
Approximate neighbors speed up KD-Tree [OK]

Hint: Approximate neighbors speed up KD-Tree [OK]

Common Mistakes:

Recommended

Ready to prove it?

Pass the test and this pattern moves unknown strong while earning +1 Grip toward your certificate.

10 questions~5 mineasy hard

Implement K-Means Clustering - Lloyd's Algorithm from Scratch

Practice

Solution

Step 1: Understand the Process

Step 2: Determine Complexity

Final Answer:

Quick Check:

Solution

Step 1: Analyze the operations

Step 2: Eliminate wrong options

Final Answer:

Quick Check:

Solution

Step 1: Identify operations

Step 2: Eliminate incorrect complexities

Final Answer:

Quick Check:

Solution

Step 1: Compare properties

Step 2: Eliminate incorrect reasons

Final Answer:

Quick Check:

Solution

Step 1: Identify optimization strategies

Step 2: Evaluate other options

Final Answer:

Quick Check:

Ready to prove it?