fix bug

apache · Oct 13, 2014 · af6dc97 · af6dc97
1 parent 2a8267a
commit af6dc97
Showing 1 changed file with 19 additions and 15 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
@@ -923,20 +923,25 @@ object DecisionTree extends Serializable with Logging {
             splits(featureIndex) = new Array[Split](numSplits)
             bins(featureIndex) = new Array[Bin](numBins)
 
-            for (splitIndex <- 0 until numSplits) {
-              val threshold = featureSplits(splitIndex)
-              splits(featureIndex)(splitIndex) =
-                new Split(featureIndex, threshold, Continuous, List())
-            }
-            bins(featureIndex)(0) = new Bin(new DummyLowSplit(featureIndex, Continuous),
-              splits(featureIndex)(0), Continuous, Double.MinValue)
-            for (splitIndex <- 1 until numSplits) {
-              bins(featureIndex)(splitIndex) =
-                new Bin(splits(featureIndex)(splitIndex - 1), splits(featureIndex)(splitIndex),
-                  Continuous, Double.MinValue)
+            if (numSplits == 0) {
+              bins(featureIndex)(0) = new Bin(new DummyLowSplit(featureIndex, Continuous),
+                new DummyLowSplit(featureIndex, Continuous), Continuous, Double.MinValue)
+            } else {
+              for (splitIndex <- 0 until numSplits) {
+                val threshold = featureSplits(splitIndex)
+                splits(featureIndex)(splitIndex) =
+                  new Split(featureIndex, threshold, Continuous, List())
+              }
+              bins(featureIndex)(0) = new Bin(new DummyLowSplit(featureIndex, Continuous),
+                splits(featureIndex)(0), Continuous, Double.MinValue)
+              for (splitIndex <- 1 until numSplits) {
+                bins(featureIndex)(splitIndex) =
+                  new Bin(splits(featureIndex)(splitIndex - 1), splits(featureIndex)(splitIndex),
+                    Continuous, Double.MinValue)
+              }
+              bins(featureIndex)(numSplits) = new Bin(splits(featureIndex)(numSplits - 1),
+                new DummyHighSplit(featureIndex, Continuous), Continuous, Double.MinValue)
             }
-            bins(featureIndex)(numSplits) = new Bin(splits(featureIndex)(numSplits - 1),
-              new DummyHighSplit(featureIndex, Continuous), Continuous, Double.MinValue)
           } else {
             // Categorical feature
             val featureArity = metadata.featureArity(featureIndex)
@@ -1050,8 +1055,7 @@ object DecisionTree extends Serializable with Logging {
         }
         index += 1
       }
-      // last value is not put into valueCount
-      // because we should not use it as a split threshold
+
       valueCount.append((currentValue, currentCount))
 
       valueCount.toArray