1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244

"""
Lloyd's Algorithm: K-Means Clustering from Scratch
 
A complete, well-documented implementation emphasizing understanding
over micro-optimizations. Each step maps directly to the theory.
"""
import numpy as np
from typing import Tuple, List, Optional
 
class KMeans:
    """
    K-Means clustering using Lloyd's algorithm.
    
    Attributes:
        k: Number of clusters
        max_iters: Maximum iterations before stopping
        tol: Convergence tolerance for centroid movement
        centroids: Final cluster centroids after fitting
        labels: Cluster assignments for training data
        inertia: Final within-cluster sum of squares
        n_iters: Number of iterations until convergence
    """
    
    def __init__(
        self,
        k: int,
        max_iters: int = 300,
        tol: float = 1e-4,
        random_state: Optional[int] = None
    ):
        self.k = k
        self.max_iters = max_iters
        self.tol = tol
        self.random_state = random_state
        
        # These are set during fit()
        self.centroids = None
        self.labels = None
        self.inertia = None
        self.n_iters = 0
        
    def _initialize_centroids(self, X: np.ndarray) -> np.ndarray:
        """
        Initialize centroids using random selection (Forgy method).
        
        Randomly selects k data points as initial centroids.
        This is simple but can lead to poor initializations.
        
        Args:
            X: Data matrix of shape (n_samples, n_features)
            
        Returns:
            Initial centroids of shape (k, n_features)
        """
        n_samples = X.shape[0]
        
        if self.random_state is not None:
            np.random.seed(self.random_state)
        
        # Select k unique indices
        indices = np.random.choice(n_samples, size=self.k, replace=False)
        
        # Return copies of selected points as centroids
        return X[indices].copy()
    
    def _compute_distances(
        self, 
        X: np.ndarray, 
        centroids: np.ndarray
    ) -> np.ndarray:
        """
        Compute squared Euclidean distances from each point to each centroid.
        
        Uses the identity ||x - μ||² = ||x||² - 2x·μ + ||μ||²
        for efficient vectorized computation.
        
        Args:
            X: Data matrix (n_samples, n_features)
            centroids: Centroid matrix (k, n_features)
            
        Returns:
            Distance matrix (n_samples, k) where D[i,j] = ||x_i - μ_j||²
        """
        # ||x||² for each sample: shape (n_samples, 1)
        X_sq = np.sum(X ** 2, axis=1, keepdims=True)
        
        # ||μ||² for each centroid: shape (1, k)
        centroids_sq = np.sum(centroids ** 2, axis=1, keepdims=True).T
        
        # -2 * X @ μ^T: shape (n_samples, k)
        cross_term = -2 * X @ centroids.T
        
        # ||x - μ||² = ||x||² - 2x·μ + ||μ||²
        distances_sq = X_sq + cross_term + centroids_sq
        
        # Clip negative values (numerical precision issues)
        return np.maximum(distances_sq, 0)
    
    def _assign_clusters(
        self, 
        X: np.ndarray, 
        centroids: np.ndarray
    ) -> np.ndarray:
        """
        Assignment step: assign each point to the nearest centroid.
        
        This creates a Voronoi tessellation of the feature space.
        
        Args:
            X: Data matrix (n_samples, n_features)
            centroids: Current centroids (k, n_features)
            
        Returns:
            Cluster labels (n_samples,) where labels[i] ∈ {0, ..., k-1}
        """
        distances_sq = self._compute_distances(X, centroids)
        return np.argmin(distances_sq, axis=1)
    
    def _update_centroids(
        self, 
        X: np.ndarray, 
        labels: np.ndarray
    ) -> np.ndarray:
        """
        Update step: recompute centroids as cluster means.
        
        For each cluster j: μ_j = (1/|C_j|) Σ_{x_i ∈ C_j} x_i
        
        Args:
            X: Data matrix (n_samples, n_features)
            labels: Current cluster assignments (n_samples,)
            
        Returns:
            New centroids (k, n_features)
        """
        n_features = X.shape[1]
        new_centroids = np.zeros((self.k, n_features))
        
        for j in range(self.k):
            # Get all points assigned to cluster j
            cluster_mask = (labels == j)
            cluster_points = X[cluster_mask]
            
            if len(cluster_points) > 0:
                # Centroid is the mean of cluster points
                new_centroids[j] = cluster_points.mean(axis=0)
            else:
                # Empty cluster: reinitialize randomly
                # This handles edge cases where a cluster loses all points
                random_idx = np.random.randint(X.shape[0])
                new_centroids[j] = X[random_idx]
                
        return new_centroids
    
    def _compute_inertia(
        self, 
        X: np.ndarray, 
        labels: np.ndarray, 
        centroids: np.ndarray
    ) -> float:
        """
        Compute within-cluster sum of squares (WCSS/inertia).
        
        Inertia = Σ_j Σ_{x_i ∈ C_j} ||x_i - μ_j||²
        
        This is the objective function that k-means minimizes.
        """
        distances_sq = self._compute_distances(X, centroids)
        # Sum of squared distances to assigned centroids
        return np.sum(distances_sq[np.arange(len(labels)), labels])
    
    def fit(self, X: np.ndarray) -> 'KMeans':
        """
        Fit k-means clustering to data using Lloyd's algorithm.
        
        The main loop alternates between:
        1. Assignment: assign points to nearest centroid
        2. Update: recompute centroids as cluster means
        
        Args:
            X: Training data (n_samples, n_features)
            
        Returns:
            self (fitted estimator)
        """
        X = np.asarray(X, dtype=np.float64)
        
        # Step 1: Initialize centroids
        self.centroids = self._initialize_centroids(X)
        
        # Main Lloyd's algorithm loop
        for iteration in range(self.max_iters):
            # Store old centroids for convergence check
            old_centroids = self.centroids.copy()
            
            # Step 2a: Assignment - assign points to nearest centroid
            self.labels = self._assign_clusters(X, self.centroids)
            
            # Step 2b: Update - recompute centroids
            self.centroids = self._update_centroids(X, self.labels)
            
            # Check convergence: have centroids stopped moving?
            centroid_shift = np.sum((self.centroids - old_centroids) ** 2)
            
            self.n_iters = iteration + 1
            
            if centroid_shift < self.tol:
                break
        
        # Compute final inertia
        self.inertia = self._compute_inertia(X, self.labels, self.centroids)
        
        return self
    
    def predict(self, X: np.ndarray) -> np.ndarray:
        """
        Predict cluster labels for new data points.
        
        Simply assigns each point to the nearest centroid.
        """
        if self.centroids is None:
            raise ValueError("Model not fitted. Call fit() first.")
        return self._assign_clusters(np.asarray(X), self.centroids)
 
 
# Example usage and visualization
if __name__ == "__main__":
    # Generate sample data: 3 clusters
    np.random.seed(42)
    
    cluster_1 = np.random.randn(50, 2) + np.array([0, 0])
    cluster_2 = np.random.randn(50, 2) + np.array([5, 5])
    cluster_3 = np.random.randn(50, 2) + np.array([10, 0])
    
    X = np.vstack([cluster_1, cluster_2, cluster_3])
    
    # Fit k-means
    kmeans = KMeans(k=3, random_state=42)
    kmeans.fit(X)
    
    print(f"Converged in {kmeans.n_iters} iterations")
    print(f"Final inertia: {kmeans.inertia:.2f}")
    print(f"Cluster sizes: {np.bincount(kmeans.labels)}")
    print(f"\nCentroids:\n{kmeans.centroids}")