From 79c35550447ccb914256c1b96702dd76ab065a53 Mon Sep 17 00:00:00 2001
From: Alexander Ulanov <nashb@yandex.ru>
Date: Tue, 15 Jul 2014 13:30:44 +0400
Subject: [PATCH] Addressing reviewers comments mengxr

---
 .../mllib/evaluation/MulticlassMetrics.scala  | 20 +++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/evaluation/MulticlassMetrics.scala b/mllib/src/main/scala/org/apache/spark/mllib/evaluation/MulticlassMetrics.scala
index d48068719a851..292cbaf639f90 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/evaluation/MulticlassMetrics.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/evaluation/MulticlassMetrics.scala
@@ -45,8 +45,9 @@ class MulticlassMetrics(predictionAndLabels: RDD[(Double, Double)]) {
       (prediction, if (prediction != label) 1 else 0)
     }.reduceByKey(_ + _)
     .collectAsMap()
-  private lazy val confusions = predictionAndLabels.map {
-    case (prediction, label) => ((prediction, label), 1)
+  private lazy val confusions = predictionAndLabels
+    .map { case (prediction, label) =>
+      ((prediction, label), 1)
   }.reduceByKey(_ + _).collectAsMap()
 
   /**
@@ -55,11 +56,18 @@ class MulticlassMetrics(predictionAndLabels: RDD[(Double, Double)]) {
    * they are ordered by class label ascending,
    * as in "labels"
    */
-  lazy val confusionMatrix: Matrix = {
+  def confusionMatrix: Matrix = {
     val transposedFlatMatrix = Array.ofDim[Double](labels.size * labels.size)
-    for (i <- 0 to labels.size - 1; j <- 0 to labels.size - 1) {
-      transposedFlatMatrix(i * labels.size + j)
-        = confusions.getOrElse((labels(i), labels(j)), 0).toDouble
+    val n = labels.size
+    var i, j = 0
+    while(i < n){
+      j = 0
+      while(j < n){
+        transposedFlatMatrix(i * labels.size + j)
+          = confusions.getOrElse((labels(i), labels(j)), 0).toDouble
+        j += 1
+      }
+      i += 1
     }
     Matrices.dense(labels.size, labels.size, transposedFlatMatrix)
   }