Quick version of k-means with numpy written in a functional style

gistfile1.py

# -*- coding: utf-8 -*-

import numpy as np
import scipy.ndimage.measurements
import random

def kMeans( data, k, centers=None, iter=64 ):
  if len(data.shape) == 1:
    data = np.vstack(data)

  if not centers:
    centers = random.sample( data, k )
  if type(centers) != np.ndarray:
    centers = np.array( centers )

  for i in range(iter):
    centers = np.apply_along_axis(
      scipy.ndimage.measurements.mean, 0, data,
      np.argmin(
        np.sum( (np.atleast_3d(data) - np.atleast_3d(centers).T)**2, axis=1 ),
        axis = 1
      ),
      range(k)
    )

  return centers