diff --git a/machine_learning/k_means_clust.py b/machine_learning/k_means_clust.py index a55153628f9c..d553d2a1e0e5 100644 --- a/machine_learning/k_means_clust.py +++ b/machine_learning/k_means_clust.py @@ -81,6 +81,13 @@ def centroid_pairwise_dist(x, centroids): def assign_clusters(data, centroids): + """Assign each data point to the index of its nearest centroid. + + >>> data = np.array([[0.0, 0.0], [0.0, 1.0], [10.0, 10.0], [10.0, 11.0]]) + >>> centroids = np.array([[0.0, 0.0], [10.0, 10.0]]) + >>> assign_clusters(data, centroids).tolist() + [0, 0, 1, 1] + """ # Compute distances between each data point and the set of centroids: # Fill in the blank (RHS only) distances_from_centroids = centroid_pairwise_dist(data, centroids) @@ -93,6 +100,13 @@ def assign_clusters(data, centroids): def revise_centroids(data, k, cluster_assignment): + """Recompute each centroid as the mean of the points assigned to it. + + >>> data = np.array([[0.0, 0.0], [0.0, 1.0], [10.0, 10.0], [10.0, 11.0]]) + >>> assignment = np.array([0, 0, 1, 1]) + >>> revise_centroids(data, 2, assignment).tolist() + [[0.0, 0.5], [10.0, 10.5]] + """ new_centroids = [] for i in range(k): # Select all data points that belong to cluster i. Fill in the blank (RHS only) @@ -106,6 +120,16 @@ def revise_centroids(data, k, cluster_assignment): def compute_heterogeneity(data, k, centroids, cluster_assignment): + """Sum of squared distances from each point to its assigned centroid. + + This is the objective k-means minimises; lower is a tighter clustering. + + >>> data = np.array([[0.0, 0.0], [0.0, 1.0], [10.0, 10.0], [10.0, 11.0]]) + >>> centroids = np.array([[0.0, 0.5], [10.0, 10.5]]) + >>> assignment = np.array([0, 0, 1, 1]) + >>> float(compute_heterogeneity(data, 2, centroids, assignment)) + 1.0 + """ heterogeneity = 0.0 for i in range(k): # Select all data points that belong to cluster i. Fill in the blank (RHS only) @@ -154,7 +178,16 @@ def kmeans( as function of iterations if None, do not store the history. verbose: if True, print how many data points changed their cluster labels in - each iteration""" + each iteration + + >>> data = np.array([[0.0, 0.0], [0.0, 1.0], [10.0, 10.0], [10.0, 11.0]]) + >>> initial_centroids = np.array([[0.0, 0.0], [10.0, 10.0]]) + >>> centroids, assignment = kmeans(data, 2, initial_centroids, maxiter=10) + >>> centroids.tolist() + [[0.0, 0.5], [10.0, 10.5]] + >>> assignment.tolist() + [0, 0, 1, 1] + """ centroids = initial_centroids[:] prev_cluster_assignment = None