Introduce Associated Labels to centroids.

This could help showing how to use KMeans as a classifier.

Some changes to KMeans properties and function.
Updated Tests accordingly.
Demo the usage of labels in one small test.

Author: Matteo Piombo

Diff for: K-Means/KMeans.swift

@@ -5,21 +5,24 @@
 
 import Foundation
 
-class KMeans {
+class KMeans<Label: Hashable> {
   let numCenters: Int
-  let convergeDist: Double
+  let labels: Array<Label>
+  private(set) var centroids: Array<Vector>
 
-  init(numCenters: Int, convergeDist: Double) {
-    self.numCenters = numCenters
-    self.convergeDist = convergeDist
+  init(labels: Array<Label>) {
+    assert(labels.count > 1, "Exception: KMeans with less than 2 centers.")
+    self.labels = labels
+    self.numCenters = labels.count
+    centroids = []
   }
 
-  private func nearestCenter(x: Vector, centers: [Vector]) -> Int {
+  private func nearestCenterIndex(x: Vector, centers: [Vector]) -> Int {
     var nearestDist = DBL_MAX
     var minIndex = 0
-
-    for (idx, c) in centers.enumerate() {
-      let dist = x.distTo(c)
+    
+    for (idx, center) in centers.enumerate() {
+      let dist = x.distTo(center)
       if dist < nearestDist {
         minIndex = idx
         nearestDist = dist
@@ -28,24 +31,26 @@ class KMeans {
     return minIndex
   }
 
-  func findCenters(points: [Vector]) -> [Vector] {
+  
+  
+  func trainCenters(points: [Vector], convergeDist: Double) {
+    
     var centerMoveDist = 0.0
-    let zeros = [Double](count: points[0].length, repeatedValue: 0.0)
+    let zeroVector = Vector(d: [Double](count: points[0].length, repeatedValue: 0.0))
 
     var kCenters = reservoirSample(points, k: numCenters)
 
     repeat {
-      var cnts = [Double](count: numCenters, repeatedValue: 0.0)
-      var newCenters = [Vector](count:numCenters, repeatedValue: Vector(d:zeros))
-
+      
+      var classification: Array<[Vector]> = Array(count: numCenters, repeatedValue: [])
+      
       for p in points {
-        let c = nearestCenter(p, centers: kCenters)
-        cnts[c] += 1
-        newCenters[c] += p
+        let classIndex = nearestCenterIndex(p, centers: kCenters)
+        classification[classIndex].append(p)
       }
 
-      for idx in 0..<numCenters {
-        newCenters[idx] /= cnts[idx]
+      let newCenters = classification.map { assignedPoints in
+        assignedPoints.reduce(zeroVector, combine: +) / Double(assignedPoints.count)
       }
 
       centerMoveDist = 0.0
@@ -56,7 +61,20 @@ class KMeans {
       kCenters = newCenters
     } while centerMoveDist > convergeDist
 
-    return kCenters
+    centroids = kCenters
+  }
+  
+  func fit(point: Vector) -> Label {
+    assert(!centroids.isEmpty, "Exception: KMeans tried to fit on a non trained model.")
+    
+    let centroidIndex = nearestCenterIndex(point, centers: centroids)
+    return labels[centroidIndex]
+  }
+  
+  func fit(points: [Vector]) -> [Label] {
+    assert(!centroids.isEmpty, "Exception: KMeans tried to fit on a non trained model.")
+    
+    return points.map(fit)
   }
 }
 

Diff for: K-Means/Tests/KMeansTests.swift

@@ -10,55 +10,66 @@ import Foundation
 import XCTest
 
 class KMeansTests: XCTestCase {
-    var points = [Vector]()
 
-    func genPoints(numPoints:Int, numDimmensions:Int) {
+    func genPoints(numPoints:Int, numDimmensions:Int) -> [Vector] {
+        var points = [Vector]()
+        
         for _ in 0..<numPoints {
             var data = [Double]()
             for _ in 0..<numDimmensions {
                 data.append(Double(arc4random_uniform(UInt32(numPoints*numDimmensions))))
             }
             points.append(Vector(d: data))
         }
-
+        return points
     }
 
     func testSmall_2D() {
-        genPoints(10, numDimmensions: 2)
-        
+        let points = genPoints(10, numDimmensions: 2)
+
         print("\nCenters")
-        let kmm = KMeans(numCenters: 3, convergeDist: 0.01)
-        for c in kmm.findCenters(points) {
-            print(c)
+        let kmm = KMeans<Character>(labels: ["A", "B", "C"])
+        kmm.trainCenters(points, convergeDist: 0.01)
+        
+        for (label, centroid) in zip(kmm.labels, kmm.centroids) {
+            print("\(label): \(centroid)")
+        }
+        
+        print("\nClassifications")
+        for (label, point) in zip(kmm.fit(points), points) {
+            print("\(label): \(point)")
         }
     }
 
     func testSmall_10D() {
-        genPoints(10, numDimmensions: 10)
+        let points = genPoints(10, numDimmensions: 10)
 
         print("\nCenters")
-        let kmm = KMeans(numCenters: 3, convergeDist: 0.01)
-        for c in kmm.findCenters(points) {
+        let kmm = KMeans<Int>(labels: [1, 2, 3])
+        kmm.trainCenters(points, convergeDist: 0.01)
+        for c in kmm.centroids {
             print(c)
         }
     }
 
     func testLarge_2D() {
-        genPoints(10000, numDimmensions: 2)
+        let points = genPoints(10000, numDimmensions: 2)
 
         print("\nCenters")
-        let kmm = KMeans(numCenters: 5, convergeDist: 0.01)
-        for c in kmm.findCenters(points) {
+        let kmm = KMeans<Character>(labels: ["A","B","C","D","E"])
+        kmm.trainCenters(points, convergeDist: 0.01)
+        for c in kmm.centroids {
             print(c)
         }
     }
 
     func testLarge_10D() {
-        genPoints(10000, numDimmensions: 10)
+        let points = genPoints(10000, numDimmensions: 10)
 
         print("\nCenters")
-        let kmm = KMeans(numCenters: 5, convergeDist: 0.01)
-        for c in kmm.findCenters(points) {
+        let kmm = KMeans<Int>(labels: [1,2,3,4,5])
+        kmm.trainCenters(points, convergeDist: 0.01)
+        for c in kmm.centroids {
             print(c)
         }
     }

Diff for: K-Means/Tests/Tests.xcodeproj/project.pbxproj

@@ -13,7 +13,7 @@
 /* End PBXBuildFile section */
 
 /* Begin PBXFileReference section */
-		B80894DB1C852CFA0018730E /* KMeans.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; name = KMeans.swift; path = ../KMeans.swift; sourceTree = SOURCE_ROOT; };
+		B80894DB1C852CFA0018730E /* KMeans.swift */ = {isa = PBXFileReference; indentWidth = 2; lastKnownFileType = sourcecode.swift; name = KMeans.swift; path = ../KMeans.swift; sourceTree = SOURCE_ROOT; tabWidth = 2; };
 		B80894E01C852D100018730E /* Tests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = Tests.xctest; sourceTree = BUILT_PRODUCTS_DIR; };
 		B80894E31C852D100018730E /* KMeansTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = KMeansTests.swift; sourceTree = SOURCE_ROOT; };
 		B80894E51C852D100018730E /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist; path = Info.plist; sourceTree = "<group>"; };