I have written a k-means function in Python to understand the methodology. I am trying to use this on a more complex dataset with a larger value for k, but it is running super slow. Does anyone have any recommendations for how I can improve this? I have provided code below, along with loading in an example dataset and applying the algorithm.

def Euc(x,y):
    return math.sqrt(sum([(a - b) ** 2 for a,b in zip(x,y)]))

def K_means(TE,k):
    Iteration = 0
    R = []
    O_a = []
    Epoch = 0
    Tol = 1
    Old_Tol = 2
    Tol_r = []
    start_time = time.time()
    mean_cl = [[random.uniform(TE.iloc[i].min(),TE.iloc[i].max()) for i in range(len(TE.columns))] for c in range(0,k)]
    for n in range(len(TE)):
        D = [Euc(TE.iloc[n].tolist(),mean_cl[c]) for c in range(0,k)]  
    while(abs(Old_Tol - Tol) > 0.005):
        Old_Tol = Tol
        Epoch = Epoch + 1
        mean_cl = [TE.iloc[[j for j, x in enumerate(O_a) if x == i]].mean() for i in range(0,k)]
        N_a = []
        for n in range(len(TE)):
            Iteration = Iteration + 1
            D = [Euc(TE.iloc[n].tolist(),mean_cl[c]) for c in range(0,k)] 
       Tol = np.mean([x != y for x,y in zip(O_a,N_a)])
        O_a = N_a
    R.append(time.time() - start_time)
    return R

def load_Pima():
    url = "http://www.stats.ox.ac.uk/pub/PRNN/pima.tr"
    Pima_training = pd.read_csv(url,sep = '\s+')
    url = "http://www.stats.ox.ac.uk/pub/PRNN/pima.te"
    Pima_testing = pd.read_csv(url,sep = '\s+')
    Pima_training = Pima_training.iloc[1:]
    Pima_testing = Pima_testing.iloc[1:]
    Pima_training.loc[:,"type"] = Pima_training.loc[:,"type"].apply(lambda x : 0 if x == 'Yes' else 1)
    Pima_testing.loc[:,"type"] = Pima_testing.loc[:,"type"].apply(lambda x : 0 if x == 'Yes' else 1)
    Features = Pima_training.loc[:,Pima_training.columns != "type"]
    Means = Features.mean()
    SDs = Features.std()
    for name in Features.columns:
        Pima_training[name] = (Pima_training[name]-Means[name])/SDs[name]
        Pima_testing[name] = (Pima_testing[name]-Means[name])/SDs[name]
    return Pima_training, Pima_testing

Pima_training, Pima_testing = load_Pima()

class_var = "type"
k = 2
TE = Pima_testing
TE = TE.loc[:,TE.columns != class_var]
km = K_means(TE,k)

The function returns the runtime of the algorithm, the tolerance at each epoch (% of changes in cluster assignment), the final cluster assignments, and the total number of iterations. I have already removed four for loops, which has sped it up quite a bit. But I fear my lack of Python programming is holding me back from making this more efficient. Any help is appreciated!

  • 1
    \$\begingroup\$ I tried to run this code, but I get this error: NameError: name 'Euc' is not defined. Can you include that function too? \$\endgroup\$
    – Nick ODell
    Commented Oct 1, 2020 at 20:10
  • 3
    \$\begingroup\$ Can you show the Euc function? One efficiency step would be to vectorize this and apply it to the whole dataframe TE at once instead of row by row. \$\endgroup\$
    – Stuart
    Commented Oct 1, 2020 at 20:10
  • \$\begingroup\$ Woops sorry guys. Added the "Euc" function now too. \$\endgroup\$ Commented Oct 1, 2020 at 20:17
  • \$\begingroup\$ Please do not edit the code after you have received a review, this violates our question and answer policy and everyone needs to be able to see the code that the answer(s) apply to. You can create a follow up question by asking a new question with the updated code and then making a link between the 2 question. \$\endgroup\$
    – pacmaninbw
    Commented Oct 4, 2020 at 11:56

1 Answer 1


You can vectorize this at various points to apply arithmetic to the whole dataframe rather than row-by-row.

def min_euclidean(df, options):
    """ Returns the index of the series in iterable options for which df - row has minimum
    Euclidean distance """
    return pd.DataFrame(((df - series) ** 2).sum(axis=1, skipna=False) for series in options).idxmin()

def k_means(TE, k):
    iteration = 0
    epoch = 0
    tol = 1
    old_tol = 2
    tols = []
    start_time = time.time()
    mean_cl = [random.uniform(TE.min(), TE.max()) for _ in range(k)]
    o_a = min_euclidean(TE, mean_cl)
    while abs(old_tol - tol) > 0.005:
        old_tol = tol
        epoch += 1
        mean_cl = [TE[o_a == i].mean() for i in range(k)]
        n_a = min_euclidean(TE, mean_cl)
        iteration += len(TE)
        tol = (o_a != n_a).mean()
        o_a = n_a
    return time.time() - start_time, tols, n_a, iteration

Note that:

  • wherever possible we work with Pandas series or dataframes instead of lists
  • I calculate mean_cl as a list of Pandas series instead of a list of lists. This way we avoid iterating through the column or row indices of TE, which is slower. There might be some way to optimize this further by storing mean_cl as a dataframe, but I can't see an obvious way of working with it that way.
  • there is no need to take the square root for the Euclidean distance - the indices associated with minimal distance will be the same working with the squared values
  • o_a and n_a are also stored as series. We can then do (o_a!= n_a).mean() to calculate the disparity between them.
  • The biggest optimization here is probably the min_euclidean function which calculates the square difference for the entire dataframe in an optimized way, rather than iterating through it.

This may be a helpful explanation of vectorization. Working with Pandas series and dataframes in an optimal often involves a slightly different way of thinking than in regular Python.

EDIT: Here is a version using numpy instead of pandas. This fully vectorizes all the operations and is substantially faster again. It uses an optimization from this StackOverflow answer.

def k_means(TE, k):
    epoch = 0
    tol = 1
    old_tol = 2
    tols = []
    start_time = time.time()
    te = np.array(TE)
    rows, columns = te.shape
    te3 = te[:, np.newaxis]  # 3d version of te for calculating euclidean more easily
    k_range = np.arange(k)[:, np.newaxis]
    mean_cl = np.random.uniform(te.min(axis=0), te.max(axis=0), size=(k, columns))
    o_a = np.nanargmin(((te3 - mean_cl) ** 2).sum(axis=2), axis=1)
    while abs(old_tol - tol) > 0.005:
        old_tol = tol
        epoch += 1
        m = o_a == k_range     # masks for each value in range 0 to k-1
        mean_cl = m.dot(te) / m.sum(1, keepdims=True)
        n_a = np.nanargmin(((te3 - mean_cl) ** 2).sum(axis=2), axis=1)
        tol = (o_a != n_a).mean()
        o_a = n_a
    return time.time() - start_time, tols, n_a, epoch * rows
  • \$\begingroup\$ Apologies there was a missing .mean() -- edited and it should return the same results now \$\endgroup\$
    – Stuart
    Commented Oct 2, 2020 at 0:54
  • \$\begingroup\$ Thanks so much for your help! \$\endgroup\$ Commented Oct 2, 2020 at 1:58
  • \$\begingroup\$ I have edited the answer as it was sometimes giving different results from the original: skipna=False is needed in the min_euclidean function for cases where there are one or more clusters to which no rows are assigned. \$\endgroup\$
    – Stuart
    Commented Oct 7, 2020 at 15:26

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.