K-Means Clustering Algorithm in Machine Learning

Data understanding is very crucial in the field of Machine Learning. With the recent trend of collecting data, every organization has started collecting some form of it. But the concern is that most of these data are unutilized and considered digital junk. Labeling this vast amount of data with the help of annotation engineers or data engineers would be time-consuming and highly expensive. We can say the quality of data is directly proportional to the cost it acquires. But we know that the budget is always limited.

If we can not spend too much on the manpower required to label the data, why not build algorithms that can provide us with information without manual intervention? In such a scenario, unsupervised learning algorithms come into existence. They help us extract meaningful information from the junk and present it in a smooth human-readable format without any supervision. The absolute difference between Supervised and Unsupervised algorithms is explained in this blog.

Key takeaways from this blog

  • What is clustering in machine learning?
  • What is meant by the K-means algorithm? 
  • What are the basic steps of K-means clustering?
  • What are the limitations of K-means?
  • What are the possible improvements in K-means?
  • How to choose the value of K in K-means?
  • What are some inherent problems with clustering algorithms?

K-means algorithm is a type of clustering algorithm; hence, before moving towards the depth of this algorithm, let's first understand the term "Clustering."

What is clustering in machine learning?

If we define this term formally: Clustering is an unsupervised learning technique that is used to find the subgroups (also known as Clusters) in the dataset.

Let's understand it in a layman's form. Suppose we have three flavored candies in a box, Salty, Sweet, and Chocolate. Earlier, all these candies were mixed in the box, and every candy had the same shape and size. Now suppose a kid asks us to give a chocolate candy. As the candies are mixed, this task will take time. We can separate these candies into three different boxes to make this easier, each corresponding to one taste. This way of separating candies is known as clustering. 

In the image below, data samples have been divided into three clusters, and each cluster has been assigned different colors. We can say that yellow squares correspond to the 1st cluster, blue corresponds to the 2nd, and red corresponds to the 3rd.

clustering in three clusters

The clustering technique is prevalent in many fields, and hence there exist many algorithms to perform it. K-means is one of them. 

What is meant by the K-means algorithm?

James MacQueen first used the term "K-means" in 1967. It is an older technique but still very popular in the data science and machine learning industries. If we define the term formally,

K-means is a simple and elegant approach which is used to partition data samples into a pre-defined “Kdistinct and non-overlapping clusters.

The value of K in the K-means algorithm depends upon the user's choice. In the image above, the user has defined the value of K = 3. Every observation from the dataset will follow two fundamental properties,

  1. Each observation belongs to at least one of the K clusters. In simple terms, there can't be any sample that is not a part of any of the clusters.
  2. No observation will belong to more than one cluster. In simple terms, one sample can not lie in two different clusters simultaneously.

Because of its simplicity, relative robustness (work with a wider variety of datasets), and "good enough" explanation, it is still considered one of the first algorithms that data analysts use to investigate any new dataset.

What are the basic steps of K-means clustering?

Let's suppose we have some data samples (X1, X2, …, Xn), and we want to divide these data samples into "K" different clusters. We can use the K-means algorithm to perform this work based on iterative logic. The complete implementation can be summarized in five steps.

Step 1

Accept inputs of data and the number of clusters to group that data. There are n data samples in our case, and we want to divide these samples into k different clusters.

Step 2

Initialize the first K clusters. There can be two ways to do that:

  • Pick first K samples, or 
  • Take Random sampling of K elements from the available dataset.

These K samples will be treated as temporary centroids (also called mean).

Step 3

Now we have (n-k) data samples left. Based on the proximity/distance to the K centroid values, each data sample from (n-k) will be assigned to only one of the K clusters. At this stage, all the samples have been assigned to some cluster. Each record is assigned to the nearest cluster using a measure of distance (e.g., Euclidean distance)

distance formula between two samples  in k-means

Step 4

We need to calculate the new centroids for every cluster. But this time, the centroid calculation will not be random, but the samples present in the cluster will be averaged to calculate the mean. At this step, we will have K centroids that can be different from the earlier chosen random K samples.

Step 5

We will iterate over this algorithm starting from Step 3 for the entire n samples. This will be repeated until we see no/negligible movement of samples among clusters. 

In this way, K-means groups the data samples into the desired number of clusters. The steps are very straightforward and very intuitive. This can be considered one of the most important reasons for the popularity of this algorithm.

What are the limitations of K-means?

K-means is a vital algorithm, but it has certain limitations as well. Some of these limitations are:

  • Looking carefully at the steps mentioned above, we can say that this algorithm is computationally expensive. It requires time proportional to the product of the number of data items (n), number of clusters (k), and the number of iterations (m).
  • In the first step, we picked K clusters randomly, because of which K-means generally converge to a local optimum. The quality of the resulting clusters highly depends upon the selection of initial K clusters.
  • There can be situations where this algorithm can suffer from an empty cluster problem. If we remember, we picked K initial values for K centroids randomly, and we may probably pick an outlier here. In such conditions, the cluster might become empty.

What are the possible improvements in K-means?

As per the limitations discussed above, there are scopes where K-means can be improved further. The K-means algorithm computes the distance between data samples and the centroid within each cluster in each iteration. It is making this algorithm computationally very expansive in the case of huge datasets. However, we can utilize the knowledge of the distance calculated in the previous iteration to make it computationally cheaper. Can you think how?

In the first iteration, we calculated the distance for each data point to the nearest cluster and termed it as previous_distance[i] for ith sample. At the next iteration, the centroids of all the clusters will change because samples will change their clusters. Here, we will first calculate the distance of the same ith sample to the previous cluster (but with the updated centroid) and termed it as current_distance[i]. Now we will compare it with the previous distance previous_distance[i]. There will be two scenarios:

  1. currentdistance[i] ≤ previousdistance[i] : There will be no movement of the ith sample within clusters. We can interpret it with the situation that the updated centroid of the cluster came closer to the ith sample, and hence it became a more potential candidate for falling into that cluster.
  2. currentdistance[i] >previousdistance[i] : Sample will leave the previous cluster and move towards the other cluster.

With this comparison at the start, we will be saving the time required to compute the distances to K-1 clusters, and it becomes slightly better in terms of computational complexity.

How to decide the number of clusters "K" in the K-means algorithm?

The decision of an optimal number of clusters is subjective and depends upon the methods to find the similarity based on which we performed the clustering. There are some algorithmic methods also that can help us in making decisions on K.

Elbow method in K-means

Let's discuss the pseudo-code using which we can implement the Elbow method.

  • Run the K-means algorithm for different values of K, let's say K=1 to K=10.
  • For every K, calculate the total sum of distances of samples to their closest cluster centroid. Let's say it SSE (Sum Squared Error). We say it is an error because ideally, every sample should lie on the centroid of the corresponding cluster, which is not possible.
  • Plot the curve of SSE with respect to the number of clusters K.
  • The value of K at the bend (knee) in the plot is generally considered an indicator of the appropriate number of clusters.
SSE = {}
for k in range(1, 10):
    k_means = KMeans(n_clusters=k, max_iter=1000).fit(data)
    data["clusters"] = kmeans.labels_
    SSE[k] = kmeans.inertia_
# Inertia: Sum of distances of samples to their closest cluster #center
plt.plot(list(SSE.keys()), list(SSE.values()))
plt.xlabel("Number of cluster")

Elbow point

We can say that the optimal value of clusters is 3 for the data used in the above code. There is another famous method,

Average silhouette method in K-means

Let's learn the pseudo-code of this algorithm as well.

  • Run the K-means algorithm for different values of K, let's say K=2 to K=11.
  • For every K, calculate the average silhouette of observations. Let's say AvS.
  • Plot the curve of AvS with respect to the number of clusters K.
  • The value of K at the maximum AvS will be considered optimal.
for n_cluster in range(2, 11):
    kmeans = KMeans(n_clusters=n_cluster).fit(data)
    label = kmeans.labels_
    AvS = silhouette_score(data, label, metric='euclidean')
    print("For n_clusters={}, The Silhouette Coefficient is {}".format(n_cluster, AvS))

For n_clusters=2, The Silhouette Coefficient is 0.680813620271
For n_clusters=3, The Silhouette Coefficient is 0.552591944521
For n_clusters=4, The Silhouette Coefficient is 0.496992849949
For n_clusters=5, The Silhouette Coefficient is 0.488517550854
For n_clusters=6, The Silhouette Coefficient is 0.370380309351
For n_clusters=7, The Silhouette Coefficient is 0.356303270516
For n_clusters=8, The Silhouette Coefficient is 0.365164535737
For n_clusters=9, The Silhouette Coefficient is 0.346583642095
For n_clusters=10, The Silhouette Coefficient is 0.328266088778


Here, n = 2 has the highest coefficient, so ideally, it should be selected, but, based on the knowledge of the Iris data, we were aware that three species are there, so the next value of 3 is chosen for optimal. When dealing with higher dimensions, this method produces a better result. From this example, we can say that these methods can only help us decide the appropriate number of clusters, and we can not wholly rely on their predictions.

Implementation of K-means in python

Let's try to implement this algorithm using the Scikit-Learn library on one of the famous datasets of the framework, i.e., the Iris Dataset shown in the image below.

Iris dataset

Step 1: Importing the required libraries of pandas, Iris datasets, KMeans model, and matplotlib.

import pandas as pd
from sklearn.datasets import load_iris
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

Step 2: Extract the data from the load_iris dataset and separate the target variables.

iris = load_iris()
X = iris.data
y = iris.target

Step 3: Fit the Scikit-learn K-means model

kmeans = KMeans(n_clusters=3, max_iter=1000).fit(X)
labels = kmeans.labels_

Step 4: Plot the scattered pred

plt.figure(figsize=(4, 3))
ax = Axes3D(fig, rect=[0, 0, .95, 1], elev=48, azim=134)

ax.scatter(X[:, 3], X[:, 0], X[:, 2],
           c=labels.astype(float), edgecolor='k')

ax.set_xlabel('Petal width')
ax.set_ylabel('Sepal length')
ax.set_zlabel('Petal length')
ax.dist = 12

3D Plot of iris data cluster

We can see that the 3 clusters are formed based on the Petal width, Petal Length, and Sepal Length. From the Average silhouette method, we received the best number of clusters is 2, and we also noticed from the above image that yellow and purple samples are very close. That's why the silhouette method got confused and suggested two clusters.

Practical Issues in Clustering

There is no doubt that clustering is one of the most important processes for data science and machine learning, but it has some practical challenges for which the solution is challenging to find. Some of these practical issues are

  1. Performing normalization or standardization before applying cluster sometimes produce better results and sometimes worse. So this decision is not firm and only be observed by actually implementing it. And the final clusters are highly dependent upon these steps.
  2. Choice of the number of clusters is not fixed and only can be observed by actually implementing it.
  3. The results of the clustering algorithms are hard to analyze in the case of higher dimensions. Dimensionality reduction methods like PCA can be used to reduce the dimensions here, but it also causes data loss.

Industrial Applications of K-means

K-means is a very popular algorithm widely used to analyze real-life data, including medical, sales, finance, and many more. Some of the most excellent applications of K-means are:

  • This algorithm is used to construct the Yolo-v3,4,5 object detection models, which are very famous in deep learning.
  • Customer Segmentation is one of the most incredible examples which uses K-means directly. More details can be found here.
  • Pattern recognition in images: We can feed millions and billions of pictures into this algorithm. K-means can segregate them into the desired number of clusters based on their content.

Possible Interview Questions

K-means is considered one of the fundamental algorithms that every interviewer expects that candidate knows about it. Data analysis and finding patterns in the massive chunk of data is the most crucial step, and knowledge of k-means suggests that the data scientist position candidates have worked on unstructured data. Possible interview questions on this topic would be,

  • What is the step-wise process in which K-means find clusters?
  • How do we decide the number of clusters?
  • How can we improve the speed of the native K-means algorithm?
  • What are the inherent problems of clustering algorithms?
  • What are the other methods used for clustering?


In this article, we discussed the K-means algorithm in a detailed fashion. We learned the pseudo-code of K-means algorithms, inherent problems associated with the K-means algorithm, and how we can improve the time performance of the K-means algorithm. We discussed methods that can help find the optimal value K, one of the most frequent questions asked in machine learning interviews. After that, we implemented the algorithm on the Iris dataset.

Enjoy Learning, Enjoy Algorithms!

Share feedback with us

More blogs to explore

Our weekly newsletter

Subscribe to get weekly content on data structure and algorithms, machine learning, system design and oops.

© 2022 Code Algorithms Pvt. Ltd.

All rights reserved.