# KMeans in the shortest and most readable format

I'm learning Python (coming from Java) so I decided to write KMeans as a practice for the language. However I want to see how could one improve the code and making it shorter and yet readable. I still find the code rather long. Also if you have comments regarding conventions or proper practices, I would really appreciate it.

import numpy as np
from numpy import linalg as LA

def main():
#data = np.arange(20).reshape( (4,5) )
file_name = "/Users/x/Desktop/tst.csv"
with open(file_name) as f:
kmeans = KMeans()
kmeans.cluster(data, 3, 200)
print kmeans.means

class KMeans:

def __init__(self):
self.means = None
self.data_assignments = None

#data is a 2D Numpy array
def cluster(self, data, numClusters, iterations):
if numClusters < 1:
raise Exception("The number of clusters should be larger than 0.")
if numClusters > data.shape[0]:
raise Exception("The number of clusters is beyond the number of rows.")

#Pick random means
randomMeansIndices = np.random.choice( range( data.shape[0] ),
size=numClusters, replace=False )

for i in randomMeansIndices:
if self.means is not None:
self.means = np.vstack( ( self.means, data[ i ] ) )
else:
self.means = data[ i ]

for iteration in xrange(iterations):
#Data assignment
self.data_assignments = {}
for row in data:
distances = {}
meanID = 0
for mean in self.means:
distances[meanID] = LA.norm(mean - row)
meanID += 1
nearestMean = min( distances, key=lambda x: distances[x] )
if nearestMean in self.data_assignments:
self.data_assignments[nearestMean] = np.vstack( (self.data_assignments[nearestMean], row) )
else:
self.data_assignments[nearestMean] = row

#Update the means
self.means = None
for mean_data_matrix in self.data_assignments.values():
if mean_data_matrix.ndim == 1:
mean_data_matrix = mean_data_matrix.reshape((1,-1)) #reshape to make it 2D
if self.means is None:
self.means = np.mean( mean_data_matrix, axis=0 ) #the first updated mean vector
else:
newMean = np.mean( mean_data_matrix, axis=0 ) #updated mean vector
self.means = np.vstack( (self.means, newMean) )#add an updated mean

main()

• From a mathematical standpoint, I don't think you should raise an exception if the number of clusters is not within the number of rows. Zero clusters would result in no clustering, while having > numRows clusters would simply result in some of the clusters being empty. Not a cause for error. Mar 6 '14 at 2:16

You use vstack a lot to build arrays row by row, with the additional complication of initializing the array variable to None, which needs a special case in your loops then. This complication could be avoided by intializing to zero-height 2D array instead, but in any case repeated vstack use is inefficient because it copies the array every time. Better options are:

• Collect rows in a list and vstack the list in the end. Good approach when you don't know the height of the array in advance.
• Best of all, create array at once by a vectorized operation.

For example this

self.means = None
for i in randomMeansIndices:
if self.means is not None:
self.means = np.vstack( ( self.means, data[ i ] ) )
else:
self.means = data[ i ]


could be done in one line, as you can index arrays with arrays:

self.means = data[randomMeansIndices]


This code could avoid the loop

distances = {}
meanID = 0
for mean in self.means:
distances[meanID] = LA.norm(mean - row)
meanID += 1
nearestMean = min( distances, key=lambda x: distances[x] )


by computing all the norms at once (NumPy 1.8 required). argmin gives the index of the minimum.

distances = LA.norm(self.means - row, axis=1)
nearestMean = np.argmin(distances)


The KMeans class has no purpose other than holding the result of clustering. The cluster method could be a stand-alone function instead and just return means, data_assignments. If you prefer to access the two items by name, use a collections.namedtuple as the return value.