[SPARK-5365][MLlib] Refactor KMeans to reduce redundant data

If a point is selected as new centers for many runs, it would collect many redundant data. This pr refactors it.

Author: Liang-Chi Hsieh <viirya@gmail.com>

Closes #4159 from viirya/small_refactor_kmeans and squashes the following commits:

25487e6 [Liang-Chi Hsieh] Refactor codes to reduce redundant data.

1 files changed, 5 insertions, 4 deletions

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/clustering/KMeans.scala b/mllib/src/main/scala/org/apache/spark/mllib/clustering/KMeans.scala
index fc46da3a93..11633e8242 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/clustering/KMeans.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/clustering/KMeans.scala
@@ -328,14 +328,15 @@ class KMeans private (
       val chosen = data.zip(costs).mapPartitionsWithIndex { (index, pointsWithCosts) =>
         val rand = new XORShiftRandom(seed ^ (step << 16) ^ index)
         pointsWithCosts.flatMap { case (p, c) =>
-          (0 until runs).filter { r =>
+          val rs = (0 until runs).filter { r =>
             rand.nextDouble() < 2.0 * c(r) * k / sumCosts(r)
-          }.map((_, p))
+          }
+          if (rs.length > 0) Some(p, rs) else None
         }
       }.collect()
       mergeNewCenters()
-      chosen.foreach { case (r, p) =>
-        newCenters(r) += p.toDense
+      chosen.foreach { case (p, rs) =>
+        rs.foreach(newCenters(_) += p.toDense)
       }
       step += 1
     }