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:
        ncols = len(f.readline().split(","))
    data = np.loadtxt(file_name, delimiter=",", usecols=range(0,ncols-1))
    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 &lt; 1:
            raise Exception("The number of clusters should be larger than 0.")
        if numClusters &gt; 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()
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.
Initialize array to right size and assign into rows.
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.