fix bug

apache · Oct 13, 2014 · 092efcb · 092efcb
1 parent f69f47f
commit 092efcb
Showing 1 changed file with 29 additions and 29 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/tree/DecisionTree.scala b/mllib/src/main/scala/org/apache/spark/mllib/tree/DecisionTree.scala
@@ -1070,37 +1070,37 @@ object DecisionTree extends Serializable with Logging {
       // just return all possible splits
       val possibleSplits = valueCount.length
       if (possibleSplits <= numSplits) {
-        return valueCount.map(_._1)
-      }
-
-      // stride between splits
-      val stride: Double = featureSamples.length.toDouble / (numSplits + 1)
-      logDebug("stride = " + stride)
-
-      // iterate `valueCount` to find splits
-      val splits = new ArrayBuffer[Double]
-      var index = 1
-      // currentCount: sum of counts of values that have been visited
-      var currentCount = valueCount(0)._2
-      // expectedCount: expected value for `currentCount`.
-      // If `currentCount` is closest value to `expectedCount`,
-      // then current value is a split threshold.
-      // After finding a split threshold, `expectedCount` is added by stride.
-      var expectedCount = stride
-      while (index < valueCount.length) {
-        // If adding count of current value to currentCount
-        // makes currentCount less close to expectedCount,
-        // previous value is a split threshold.
-        if (math.abs(currentCount - expectedCount) <
-          math.abs(currentCount + valueCount(index)._2 - expectedCount)) {
-          splits.append(valueCount(index-1)._1)
-          expectedCount += stride
+        valueCount.map(_._1)
+      } else {
+        // stride between splits
+        val stride: Double = featureSamples.length.toDouble / (numSplits + 1)
+        logDebug("stride = " + stride)
+
+        // iterate `valueCount` to find splits
+        val splits = new ArrayBuffer[Double]
+        var index = 1
+        // currentCount: sum of counts of values that have been visited
+        var currentCount = valueCount(0)._2
+        // expectedCount: expected value for `currentCount`.
+        // If `currentCount` is closest value to `expectedCount`,
+        // then current value is a split threshold.
+        // After finding a split threshold, `expectedCount` is added by stride.
+        var expectedCount = stride
+        while (index < valueCount.length) {
+          // If adding count of current value to currentCount
+          // makes currentCount less close to expectedCount,
+          // previous value is a split threshold.
+          if (math.abs(currentCount - expectedCount) <
+            math.abs(currentCount + valueCount(index)._2 - expectedCount)) {
+            splits.append(valueCount(index-1)._1)
+            expectedCount += stride
+          }
+          currentCount += valueCount(index)._2
+          index += 1
         }
-        currentCount += valueCount(index)._2
-        index += 1
-      }
 
-      splits.toArray
+        splits.toArray
+      }
     }
 
     assert(splits.length > 0)