| 5 | return np.sqrt(np.sum((x1-x2)**2)) |
| 6 | |
| 7 | class KMeans: |
| 8 | |
| 9 | def __init__(self, K=5, max_iters=100, plot_steps=False): |
| 10 | self.K = K |
| 11 | self.max_iters = max_iters |
| 12 | self.plot_steps = plot_steps |
| 13 | |
| 14 | # list of sample indices for each cluster |
| 15 | self.clusters = [[] for _ in range(self.K)] |
| 16 | |
| 17 | # the centers (mean vector) for each cluster |
| 18 | self.centroids = [] |
| 19 | |
| 20 | |
| 21 | def predict(self, X): |
| 22 | self.X = X |
| 23 | self.n_samples, self.n_features = X.shape |
| 24 | |
| 25 | # initialize |
| 26 | random_sample_idxs = np.random.choice(self.n_samples, self.K, replace=False) |
| 27 | self.centroids = [self.X[idx] for idx in random_sample_idxs] |
| 28 | |
| 29 | # optimize clusters |
| 30 | for _ in range(self.max_iters): |
| 31 | # assign samples to closest centroids (create clusters) |
| 32 | self.clusters = self._create_clusters(self.centroids) |
| 33 | |
| 34 | if self.plot_steps: |
| 35 | self.plot() |
| 36 | |
| 37 | # calculate new centroids from the clusters |
| 38 | centroids_old = self.centroids |
| 39 | self.centroids = self._get_centroids(self.clusters) |
| 40 | |
| 41 | if self._is_converged(centroids_old, self.centroids): |
| 42 | break |
| 43 | |
| 44 | if self.plot_steps: |
| 45 | self.plot() |
| 46 | |
| 47 | # classify samples as the index of their clusters |
| 48 | return self._get_cluster_labels(self.clusters) |
| 49 | |
| 50 | |
| 51 | def _get_cluster_labels(self, clusters): |
| 52 | # each sample will get the label of the cluster it was assigned to |
| 53 | labels = np.empty(self.n_samples) |
| 54 | for cluster_idx, cluster in enumerate(clusters): |
| 55 | for sample_idx in cluster: |
| 56 | labels[sample_idx] = cluster_idx |
| 57 | |
| 58 | return labels |
| 59 | |
| 60 | |
| 61 | def _create_clusters(self, centroids): |
| 62 | # assign the samples to the closest centroids |
| 63 | clusters = [[] for _ in range(self.K)] |
| 64 | for idx, sample in enumerate(self.X): |