fixed classification requirements

markhamstra · May 13, 2014 · ed5a2df · ed5a2df
1 parent d012be7
commit ed5a2df
Showing 1 changed file with 2 additions and 2 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
@@ -1093,9 +1093,9 @@ object DecisionTree extends Serializable with Logging {
      */
     if (strategy.categoricalFeaturesInfo.size > 0) {
       val maxCategoriesForFeatures = strategy.categoricalFeaturesInfo.maxBy(_._2)._2
-      require(numBins >= maxCategoriesForFeatures)
+      require(numBins > maxCategoriesForFeatures)
       if (strategy.isMultiClassification) {
-        require(numBins > math.pow(2, maxCategoriesForFeatures.toInt) - 1)
+        require(numBins > math.pow(2, maxCategoriesForFeatures.toInt - 1) - 1)
       }
     }