Updated PR per comments from mengxr

MediaV · Jul 17, 2014 · 4e9bd1e · 4e9bd1e
1 parent 6c7a2ec
commit 4e9bd1e
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 11 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/clustering/LocalKMeans.scala b/mllib/src/main/scala/org/apache/spark/mllib/clustering/LocalKMeans.scala
@@ -62,9 +62,10 @@ private[mllib] object LocalKMeans extends Logging {
       if (j == 0) {
         logWarning("kMeansPlusPlus initialization ran out of distinct points for centers." +
           s" Using duplicate point for center k = $i.")
-        j = 1
+        centers(i) = points(0).toDense
+      } else {
+        centers(i) = points(j - 1).toDense
       }
-      centers(i) = points(j-1).toDense
     }
 
     // Run up to maxIterations iterations of Lloyd's algorithm

diff --git a/mllib/src/test/scala/org/apache/spark/mllib/clustering/KMeansSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/clustering/KMeansSuite.scala
@@ -62,11 +62,12 @@ class KMeansSuite extends FunSuite with LocalSparkContext {
   }
 
   test("no distinct points") {
-    val data = sc.parallelize(Array(
-      Vectors.dense(1.0, 2.0, 3.0),
-      Vectors.dense(1.0, 2.0, 3.0),
-      Vectors.dense(1.0, 2.0, 3.0)
-    ))
+    val data = sc.parallelize(
+      Array(
+        Vectors.dense(1.0, 2.0, 3.0),
+        Vectors.dense(1.0, 2.0, 3.0),
+        Vectors.dense(1.0, 2.0, 3.0)),
+      2)
     val center = Vectors.dense(1.0, 2.0, 3.0)
 
     // Make sure code runs.
@@ -75,10 +76,11 @@ class KMeansSuite extends FunSuite with LocalSparkContext {
   }
 
   test("more clusters than points") {
-    val data = sc.parallelize(Array(
-      Vectors.dense(1.0, 2.0, 3.0),
-      Vectors.dense(1.0, 3.0, 4.0)
-    ))
+    val data = sc.parallelize(
+      Array(
+        Vectors.dense(1.0, 2.0, 3.0),
+        Vectors.dense(1.0, 3.0, 4.0)),
+      2)
 
     // Make sure code runs.
     var model = KMeans.train(data, k=3, maxIterations=1)