StatRed ass4: Implemented part2 (k-means++).

statred/ass4/k-means.py

 from pylab import loadtxt, array, scatter, figure, show, mean, argmin, \
-        ones, append, savefig
-from random import random, seed
+        append, savefig
+from random import random, seed, randint
 from sys import argv, exit
 
 def init(X, k):
-    return X[:k]
+    """Simply use k random datapoints as initial means for the clusters."""
+    return X[[int(X.shape[0] * random()) for i in range(k)]]
 
 def init_pp(X, k):
-    return X[:k]
-
-if len(argv) == 3:
-    if argv[2] != 'pp':
-        print 'Usage: python %s K [ "pp" ]' % argv[0]
-        exit()
-    print 'Using k-means++'
-    initial_means = init_pp
-else:
+    """Use the k-means++ algorithm to find initial means."""
+    # Choose first center at random
+    N = X.shape[0]
+    indices = [int(N * random())]
+    m = [X[indices[0]]]
+    D = [((x - m[0])**2).sum() for x in X]
+    while len(m) < k:
+        best_sum = new = -1
+        for i in range(N):
+            if i not in indices:
+                Dsum = sum([min(D[j], ((X[j] - X[i])**2).sum()) for j in range(N)])
+                if best_sum == -1 or best_sum < Dsum:
+                    best_sum, new = Dsum, i
+        m.append(X[new])
+        indices.append(new)
+        D = [min(D[i], ((X[i] - X[new])**2).sum()) for i in range(N)]
+    return array(m)
+
+# Parse parameters (600 is a nice example seed)
+pp = False
+if len(argv) > 2:
+    if argv[2] == 'pp':
+        print 'Using k-means++'
+        initial_means = init_pp
+        pp = True
+    else:
+        seed(int(argv[2]))
+    if len(argv) == 4:
+        seed(int(argv[3]))
+elif len(argv) < 2:
+    print 'Usage: python %s K [ "pp" ] [ SEED ]' % argv[0]
+    exit()
+if not pp:
     print 'Using normal k-means'
     initial_means = init
 k = int(argv[1])
+if k < 1 or k > 6:
+    print 'K must be a value from 1-6'
+    exit()
 
-# Generate dataset
-seed(700)
+# Generate dataset, add a multiplication of k so that clusters are formed
 n, N = 2, 100
-X = array([[100 * random() for j in range(n)] for i in range(int(N / k + N % k))])
+X = array([[100 * random() + 70 for j in range(n)] for i in \
+        range(int(N / k + N % k))])
 for c in range(k - 1):
-    d = (k + 1) * 100 * random()
+    d = (c + 2) * 70
     X = append(X, [[100 * random() + d for j in range(n)] for i in \
         range(int(N / k))], 0)
 
+# Divide in clusters by applying k-means
 M = initial_means(X, k)
 Mp = M - 1
 steps = 0
 
-# Divide in clusters
 while (Mp - M).any():
     Mp = M
     clusters = [[] for i in range(k)]
@@ -48,7 +76,7 @@ print 'Completed in %d steps' % steps
 
 # Plot clusters
 figure(1)
-colors = [[1,0,0], [0,1,0], [0,0,1]]
+colors = [[1,0,0], [0,1,0], [0,0,1], [0,0,0], [1,1,1], [.5,.5,.5]]
 for i in range(k):
     c = array(clusters[i])
     scatter(c[:,0], c[:,1], c=colors[i])